diff --git a/LICENSE.txt b/LICENSE.txt deleted file mode 100644 index f85e365d..00000000 --- a/LICENSE.txt +++ /dev/null @@ -1,17 +0,0 @@ -Permission is hereby granted, free of charge, to any person obtaining a copy of -this software and associated documentation files (the "Software"), to deal in -the Software without restriction, including without limitation the rights to -use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies -of the Software, and to permit persons to whom the Software is furnished to do -so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in index 346766ba..1e3be698 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -5,3 +5,6 @@ include *.md # Include the license file include LICENSE + +# Include the data files +include xwakes/wit/materials.json diff --git a/README.md b/README.md index 9b2cafce..1f3514d6 100644 --- a/README.md +++ b/README.md @@ -1,55 +1,31 @@ -# PyWIT -####Python Wake and Impedance Toolbox - -## Installation -**NB:** Though most of this PyWIT's functionality is available independently of any non-dependency package, -the installation of IW2D (https://gitlab.cern.ch/IRIS/IW2D) is *strongly recommended* in order to have full access to -all of the features of PyWIT. For this reason, this installation guide will include the installation process required by -IW2D. If you already have IW2D downloaded and installed on your system, you can skip directly to step 3. Please note that IW2D (at the time of writing) is only available on Linux systems. - -1. Download the IW2D repository to your system by cloning the git repository. This can be done by executing the command - - `git clone https://gitlab.cern.ch/IRIS/IW2D` - - in a suitable directory. You may be prompted to enter your CERN log-in credentials for this step. - - -2. Perform the installation procedure for IW2D. This is explained in detail in the `README.md` file in the IW2D repository, but is summarized here: - 1. Install the prerequisite libraries "MPFR" and "GSL" by running: - `sudo apt-get install libgsl-dev libmpfr-dev` - or alternatively, if you have anaconda/miniconda installed: - `conda install -c conda-forge gsl mpfr` - 2. Compile the C++ code by running: - ``` - cd IW2D/cpp - cp ./Makefile_system_GMP_MPFR Makefile - make - ``` +# Xwakes +Python package for wakes and impedances handling. -3. In a new directory, where you want your PyWIT directory to be placed, Download the PyWIT repository to your system by - cloning the git repository. This can be done by executing the command - - `git clone https://gitlab.cern.ch/IRIS/pywit`. - - You may be prompted to enter your CERN log-in credentials for this step. - +## Installation -4. Install PyWIT using pip by navigating to your pywit directory executing the command +Under a conda environment with Python 3.8+ it is simply installed via PyPI by doing `pip install xwakes` - `pip install .`. +## IW2D coupling +This section describes how to couple Xwakes to IW2D using the executables obtained compiling the C++ code. +When the Python interface of IW2D will be completed this will not be needed anymore. -5. Navigate to the directory `IW2D/IW2D/cpp` in the IW2D repository. Copy the following four -files from this folder: +To begin with, some folders need to be created in the user's home directory. +This can be automatically done by running the following command after Xwakes is installed: +``` +python -c 'import xwakes; xwakes.initialize_pywit_directory()' +``` +The IW2D executable are produced by following the [IW2D readme]([https://gitlab.cern.ch/IRIS/IW2D/-/tree/master?ref_type=heads)](https://gitlab.cern.ch/IRIS/IW2D/). +After this procedure is completed the following executable files will be created in `/path/to/iw2d/IW2D/cpp/`: * `flatchamber.x` * `roundchamber.x` * `wake_flatchamber.x` * `wake_roundchamber.x` - - -6. In your home directory, there will have appeared a new directory called 'pywit'. Navigate into this folder, then -into 'IW2D', then into 'bin'. Paste the four files copied in step 5 into this folder. +These files have to be copied in the newly created folder with the command -You should now be able to use PyWIT with your Python system interpreter. +``` +cp /path/to/iw2d/IW2D/cpp/*.x ~/pywit/IW2D/bin +``` +Now Xwakes can be used to launch IW2D calculations. diff --git a/pywit/__init__.py b/pywit/__init__.py index e69de29b..8262aed3 100644 --- a/pywit/__init__.py +++ b/pywit/__init__.py @@ -0,0 +1 @@ +from xwakes.wit import * \ No newline at end of file diff --git a/pywit/_generate_modules.py b/pywit/_generate_modules.py new file mode 100644 index 00000000..4c095198 --- /dev/null +++ b/pywit/_generate_modules.py @@ -0,0 +1,14 @@ +''' +generate files to mirror in pywit the modules from xwakes.wit +''' + +import os + +fff = os.listdir('../xwakes/wit') + +for nn in fff: + if not nn.endswith('.py') or nn.startswith('_'): + continue + + with open(nn, 'w') as fid: + fid.write(f'from xwakes.wit.{nn.split(".py")[0]} import *\n') \ No newline at end of file diff --git a/pywit/component.py b/pywit/component.py index 9f589a92..494ae52f 100644 --- a/pywit/component.py +++ b/pywit/component.py @@ -1,386 +1 @@ -from __future__ import annotations - -from pywit.parameters import * -from pywit.utils import unique_sigfigs - -from typing import Optional, Callable, Tuple, Union, List - -import numpy as np -import sortednp as snp - - -def mix_fine_and_rough_sampling(start: float, stop: float, rough_points: int, - fine_points: int, rois: List[Tuple[float, float]]): - """ - Mix a fine and rough (geometric) sampling between start and stop, - refined in the regions of interest rois. - :param start: The lowest bound of the sampling - :param stop: The highest bound of the sampling - :param rough_points: The total number of data points to be - generated for the rough grid, between start - and stop. - :param fine_points: The number of points in the fine grid of each roi. - :param rois: List of unique tuples with the lower and upper bound of - of each region of interest. - :return: An array with the sampling obtained. - """ - intervals = [np.linspace(max(i, start), min(f, stop), fine_points) - for i, f in rois - if (start <= i <= stop or start <= f <= stop)] - fine_sampling_rois = np.hstack(intervals) if intervals else np.array([]) - rough_sampling = np.geomspace(start, stop, rough_points) - - return unique_sigfigs(snp.merge(fine_sampling_rois, rough_sampling), 7) - - -class Component: - """ - A data structure representing the impedance- and wake functions of some Component in a specified plane. - """ - - def __init__(self, impedance: Optional[Callable] = None, wake: Optional[Callable] = None, plane: str = '', - source_exponents: Tuple[int, int] = (-1, -1), test_exponents: Tuple[int, int] = (-1, -1), - name: str = "Unnamed Component", f_rois: Optional[List[Tuple[float, float]]] = None, - t_rois: Optional[List[Tuple[float, float]]] = None): - """ - The initialization function for the Component class. - :param impedance: A callable function representing the impedance function of the Component. Can be undefined if - the wake function is defined. - :param wake: A callable function representing the wake function of the Component. Can be undefined if - the impedance function is defined. - :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be specified for valid initialization - :param source_exponents: The exponents in the x and y planes experienced by the source particle. Also - referred to as 'a' and 'b'. Must be specified for valid initialization - :param test_exponents: The exponents in the x and y planes experienced by the source particle. Also - referred to as 'a' and 'b'. Must be specified for valid initialization - :param name: An optional user-specified name of the component - """ - # Enforces that either impedance or wake is defined. - assert impedance or wake, "The impedance- and wake functions cannot both be undefined." - # The impedance- and wake functions as callable objects, e.g lambda functions - self.impedance = impedance - self.wake = wake - self.name = name - - # The plane of the Component, either 'x', 'y' or 'z' - assert plane.lower() in ['x', 'y', 'z'], \ - "Cannot initialize Component object without specified plane" - self.plane = plane - - assert source_exponents != (-1, -1) and len(source_exponents) == 2, \ - "Cannot initialize Component object without specified source exponents (a, b)" - self.source_exponents = source_exponents - assert test_exponents != (-1, -1) and len(test_exponents) == 2, \ - "Cannot initialize Component object without specified test exponents (c, d)" - self.test_exponents = test_exponents - self.power_x = (source_exponents[0] + test_exponents[0] + (plane == 'x')) / 2 - self.power_y = (source_exponents[1] + test_exponents[1] + (plane == 'y')) / 2 - self.f_rois = f_rois if f_rois else [] - self.t_rois = t_rois if t_rois else [] - - def generate_wake_from_impedance(self) -> None: - """ - Uses the impedance function of the Component object to generate its wake function, using - a Fourier transform. - :return: Nothing - """ - # # If the object already has a wake function, there is no need to generate it. - # if self.wake: - # pass - # # In order to generate a wake function, we need to make sure the impedance function is defined. - # assert self.impedance, "Tried to generate wake from impedance, but impedance is not defined." - # - # raise NotImplementedError - - # Temporary solution to avoid crashes - self.wake = lambda x: 0 - - def generate_impedance_from_wake(self) -> None: - """ - Uses the wake function of the Component object to generate its impedance function, using - a Fourier transform. - :return: Nothing - """ - # # If the object already has an impedance function, there is no need to generate it. - # if self.impedance: - # pass - # # In order to generate an impedance function, we need to make sure the wake function is defined. - # assert self.wake, "Tried to generate impedance from wake, but wake is not defined." - # - # raise NotImplementedError - - # Temporary solution to avoid crashes - self.impedance = lambda x: 0 - - def is_compatible(self, other: Component) -> bool: - """ - Compares all parameters of the self-object with the argument of the function and returns True if all of their - attributes, apart from the impedance and wake functions, are identical. - i.e. Returns True if self and other are compatible for Component addition, False otherwise - :param other: Another Component - :return: True if self and other can be added together, False otherwise - """ - return all([self.source_exponents == other.source_exponents, - self.test_exponents == other.test_exponents, - self.plane == other.plane, - self.power_x == other.power_x, - self.power_y == other.power_y]) - - def __add__(self, other: Component) -> Component: - """ - Defines the addition operator for two Components - :param self: The left addend - :param other: The right addend - :return: A new Component whose impedance and wake functions are the sums - of the respective functions of the two addends. - """ - # Enforce that the two addends are in the same plane - assert self.plane == other.plane, "The two addends correspond to different planes and cannot be added.\n" \ - f"{self.plane} != {other.plane}" - - # Enforce that the two addends have the same exponent parameters - assert self.source_exponents == other.source_exponents and self.test_exponents == other.test_exponents, \ - "The two addends have different exponent parameters and cannot be added.\n" \ - f"Source: {self.source_exponents} != {other.source_exponents}\n" \ - f"Test: {self.test_exponents} != {other.test_exponents}" - - # Defines an empty array to hold the two summed functions - sums = [] - - # Iterates through the two pairs of functions: impedances, then wakes - for left, right in zip((self.impedance, self.wake), (other.impedance, other.wake)): - # If neither addend has a defined function, we will not bother to calculate that of the sum - if (not left) and (not right): - sums.append(None) - else: - # # Generates the missing function for the addend which is missing it - # if not left: - # [self.generate_impedance_from_wake, self.generate_wake_from_impedance][len(sums)]() - # elif not right: - # [other.generate_impedance_from_wake, other.generate_wake_from_impedance][len(sums)]() - # - - # TODO: Temporary fix until Fourier transform implemented - if not left: - sums.append(right) - elif not right: - sums.append(left) - else: - # Appends the sum of the functions of the two addends to the list "sums" - sums.append(lambda x, l=left, r=right: l(x) + r(x)) - - # Initializes and returns a new Component - return Component(sums[0], sums[1], self.plane, self.source_exponents, self.test_exponents, - f_rois=self.f_rois + other.f_rois, t_rois=self.t_rois + other.t_rois) - - def __radd__(self, other: Union[int, Component]) -> Component: - """ - Implements the __radd__ method for the Component class. This is only done to facilitate the syntactically - practical use of the sum() method for Components. sum(iterable) works by adding all of the elements of the - iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case - that the left addend of any addition involving a Component is not itself a Component, the resulting sum - is simply defined to be the right addend. - :param other: The left addend of an addition - :return: The sum of self and other if other is a Component, otherwise just self. - """ - # Checks if the left addend, other, is not a Component - if not isinstance(other, Component): - # In which case, the right addend is simply returned - return self - - # Otherwise, their sum is returned (by invocation of Component.__add__(self, other)) - return self + other - - def __mul__(self, scalar: complex) -> Component: - """ - Defines the behavior of multiplication of a Component by some, potentially complex, scalar - :param scalar: A scalar value to be multiplied with some Component - :return: A newly initialized Component identical to self in every way apart from the impedance- - and wake functions, which have been multiplied by the scalar. - """ - # Throws an AssertionError if scalar is not of the type complex, float or int - assert isinstance(scalar, complex) or isinstance(scalar, float) or isinstance(scalar, int) - # Initializes and returns a new Component with attributes like self, apart from the scaled functions - return Component((lambda x: scalar * self.impedance(x)) if self.impedance else None, - (lambda x: scalar * self.wake(x)) if self.wake else None, self.plane, - self.source_exponents, self.test_exponents, self.name, self.f_rois, self.t_rois) - - def __rmul__(self, scalar: complex) -> Component: - """ - Generalizes scalar multiplication of Component to be possibly from left and right. Both of these operations - are identical. - :param scalar: A scalar value to be multiplied with some Component - :return: The result of calling Component.__mul__(self, scalar): A newly initialized Component identical to self - in every way apart from the impedance- and wake functions, which have been multiplied by the scalar. - """ - # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function - return self * scalar - - def __truediv__(self, scalar: complex) -> Component: - """ - Implements the __truediv__ method for the Component class in order to produce the expected behavior when - dividing a Component by some scalar. That is, the scalar multiplication of the multiplicative inverse of - the scalar. - :param scalar: A scalar value for the Component to be divided by - :return: The result of calling Component.__mul__(self, 1 / scalar): A newly initialized Component identical to - self in every way apart from the impedance- and wake functions, which have been multiplied by (1 / scalar). - """ - # Defines the operation c / z to be equivalent to c * (1 / z) for some Component c and scalar z. - return self * (1 / scalar) - - def __str__(self) -> str: - """ - Implements the __str__ method for the Component class, providing an informative printout of the attributes - of the Component. - :return: A multi-line string containing information about the attributes of self, including which of the - impedance- and wake functions are defined. - """ - return f"{self.name} with parameters:\nPlane:\t\t\t{self.plane}\n" \ - f"Source exponents:\t{', '.join(str(i) for i in self.source_exponents)}\n" \ - f"Test exponents:\t\t{', '.join(str(i) for i in self.test_exponents)}\n" \ - f"Impedance function:\t{'DEFINED' if self.impedance else 'UNDEFINED'}\n" \ - f"Wake function:\t\t{'DEFINED' if self.wake else 'UNDEFINED'}\n" \ - f"Impedance-regions of interest: {', '.join(str(x) for x in self.f_rois)}\n" \ - f"Wake-regions of interest: {', '.join(str(x) for x in self.t_rois)}" - - def __lt__(self, other: Component) -> bool: - """ - Implements the __lt__ (less than) method for the Component class. This has no real physical interpretation, - it is just syntactically practical when iterating through the Components when adding two Elements. - (see Element.__add__) - :param other: The right hand side of a "<"-inequality - :return: True if the self-Component is "less than" the other-Component by some arbitrary, but consistent, - comparison. - """ - # The two Components are compared by their attributes - return [self.plane, self.source_exponents, self.test_exponents] < \ - [other.plane, other.source_exponents, other.test_exponents] - - def __eq__(self, other: Component) -> bool: - """ - Implements the __eq__ method for the Component class. Two Components are designated as "equal" if the following - two conditions hold: - 1. They are compatible for addition. That is, all of their attributes, apart from impedance- and wake functions, - are exactly equal. - 2. The resulting values from evaluating the impedance functions of the two components respectively for 50 - points between 1 and 10000 all need to be close within some given tolerance. The same has to hold for the wake - function. - This somewhat approximated numerical approach to the equality comparator aims to compensate for small - numerical/precision errors accumulated for two Components which have taken different "paths" to what should - analytically be identical Components. - :param other: The right hand side of the equality comparator - :return: True if the two Components have identical attributes and sufficiently close functions, False otherwise - """ - # First checks if the two Components are compatible for addition, i.e. if they have the same non-function - # attributes - if not self.is_compatible(other): - # If they do not, they are not equal - return False - - # Creates a numpy array of 50 points and verifies that the evaluations of the functions of the two components - # for all of these components are sufficiently close. If they are, True is returned, otherwise, False is - # returned. - xs = np.linspace(1, 10000, 50) - return (np.allclose(self.impedance(xs), other.impedance(xs), rtol=REL_TOL, atol=ABS_TOL) and - np.allclose(self.wake(xs), other.wake(xs), rtol=REL_TOL, atol=ABS_TOL)) - - def impedance_to_array(self, rough_points: int, start: float = MIN_FREQ, - stop: float = MAX_FREQ, - precision_factor: float = FREQ_P_FACTOR) -> Tuple[np.ndarray, np.ndarray]: - """ - Produces a frequency grid based on the f_rois attribute of the component and evaluates the component's - impedance function at these frequencies. - :param rough_points: The total number of data points to be - generated for the rough grid, between start - and stop. - :param start: The lowest frequency in the desired frequency grid - :param stop: The highest frequency in the desired frequency grid - :param precision_factor: A number indicating the ratio of points - which should be placed within the regions of interest. - If =0, the frequency grid will ignore the intervals in f_rois. - If =1, the points will be distributed 1:1 between the - rough grid and the fine grid on each roi. In general, =n - means that there will be n times more points in the fine - grid of each roi, than in the rough grid. - :return: A tuple of two numpy arrays with same shape, giving - the frequency grid and impedances respectively - """ - if len(self.f_rois) == 0: - xs = np.geomspace(start, stop, rough_points) - return xs, self.impedance(xs) - - # eliminate duplicates - f_rois_no_dup = set(self.f_rois) - - fine_points_per_roi = int(round(rough_points * precision_factor)) - - xs = mix_fine_and_rough_sampling(start, stop, rough_points, - fine_points_per_roi, - list(f_rois_no_dup)) - - return xs, self.impedance(xs) - - def wake_to_array(self, rough_points: int, start: float = MIN_TIME, - stop: float = MAX_TIME, - precision_factor: float = TIME_P_FACTOR) -> Tuple[np.ndarray, np.ndarray]: - """ - Produces a time grid based on the t_rois attribute of the component and evaluates the component's - wake function at these time points. - :param rough_points: The total number of data points to be - generated for the rough grid, between start - and stop. - :param start: The lowest time in the desired time grid - :param stop: The highest time in the desired time grid - :param precision_factor: A number indicating the ratio of points - which should be placed within the regions of interest. - If =0, the frequency grid will ignore the intervals in t_rois. - If =1, the points will be distributed 1:1 between the - rough grid and the fine grid on each roi. In general, =n - means that there will be n times more points in the fine - grid of each roi, than in the rough grid. - :return: A tuple of two numpy arrays with same shape, giving the - time grid and wakes respectively - """ - if len(self.t_rois) == 0: - xs = np.geomspace(start, stop, rough_points) - return xs, self.wake(xs) - - # eliminate duplicates - t_rois_no_dup = set(self.t_rois) - - fine_points_per_roi = int(round(rough_points * precision_factor)) - - xs = mix_fine_and_rough_sampling(start, stop, rough_points, - fine_points_per_roi, - list(t_rois_no_dup)) - - return xs, self.wake(xs) - - def discretize(self, freq_points: int, time_points: int, freq_start: float = MIN_FREQ, freq_stop: float = MAX_FREQ, - time_start: float = MIN_TIME, time_stop: float = MAX_TIME, - freq_precision_factor: float = FREQ_P_FACTOR, - time_precision_factor: float = TIME_P_FACTOR) -> Tuple[ - Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: - """ - Combines the two discretization-functions in order to fully discretize the wake and impedance of the object - as specified by a number of parameters. - :param freq_points: The total number of frequency/impedance points - :param time_points: The total number of time/wake points - :param freq_start: The lowest frequency in the frequency grid - :param freq_stop: The highest frequency in the frequency grid - :param time_start: The lowest time in the time grid - :param time_stop: The highest time in the time grid - :param freq_precision_factor: The ratio of points in the fine frequency grid - (#points in frequency ROIs / rough frequency points) - :param time_precision_factor: The ratio of points in the fine time grid - (#points in time ROIs / rough time points) - :return: A tuple of two tuples, containing the results of impedance_to_array and wake_to_array respectively - """ - return (self.impedance_to_array(freq_points, freq_start, freq_stop, freq_precision_factor), - self.wake_to_array(time_points, time_start, time_stop, time_precision_factor)) - - def get_shorthand_type(self) -> str: - """ - :return: A 5-character string indicating the plane as well as source- and test exponents of the component - """ - return self.plane + "".join(str(x) for x in (self.source_exponents + self.test_exponents)) +from xwakes.wit.component import * diff --git a/pywit/devices.py b/pywit/devices.py index 4ff77cc8..81f8db52 100644 --- a/pywit/devices.py +++ b/pywit/devices.py @@ -1,312 +1 @@ -from typing import Tuple,Sequence -from pywit.component import Component -from pywit.element import Element -from pywit.utilities import create_resonator_component -from pywit.interface import component_names, get_component_name - -import numpy as np -from scipy import integrate - -from scipy.constants import c as c_light, mu_0 -from scipy.special import erfc, zeta - -free_space_impedance = mu_0 * c_light - - -def create_tesla_cavity_component(plane: str, exponents: Tuple[int, int, int, int], a: float, g: float, - period_length: float) -> Component: - """ - Creates a single component object modeling a periodic accelerating stucture. - Follow K. Bane formalism developped in SLAC-PUB-9663, "Short-range dipole wakefields - in accelerating structures for the NLC". - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param a: accelerating structure iris gap in m - :param g: individual cell gap in m - :param period_length: period length in m - :return: A component object of a periodic accelerating structure - """ - - # Material properties required for the skin depth computation are derived from the input Layer attributes - # material_resistivity = layer.dc_resistivity - # material_relative_permeability = layer.magnetic_susceptibility - # material_permeability = material_relative_permeability * scipy.constants.mu_0 - - # Create the skin depth as a function of frequency and layer properties - # delta_skin = lambda freq: (material_resistivity/ (2*pi*abs(freq) * material_permeability)) ** (1/2) - - gamma = g / period_length - alpha1 = 0.4648 - alpha = 1 - alpha1 * np.sqrt(gamma) - (1 - 2*alpha1) * gamma - - s00 = g / 8 * (a / (alpha * period_length)) ** 2 - - # Longitudinal impedance and wake - if plane == 'z' and exponents == (0, 0, 0, 0): - def longitudinal_impedance_tesla_cavity(freq): - return (1j * free_space_impedance / (np.pi * (2 * np.pi * freq / c_light) * a ** 2) * - (1 + (1 + 1j) * alpha * period_length / a * np.sqrt(np.pi / ((2 * np.pi * freq / c_light) * g))) ** (-1)) - - def longitudinal_wake_tesla_cavity(time): - return ((free_space_impedance * c_light / (np.pi * a ** 2)) * np.heaviside(time, 0) * - np.exp(-np.pi * c_light * time / (4 * s00)) * erfc(np.sqrt(np.pi * c_light * time / (4 * s00)))) - - impedance = longitudinal_impedance_tesla_cavity - wake = longitudinal_wake_tesla_cavity - # Transverse dipolar impedance and wake - elif (plane == 'x' and exponents == (1, 0, 0, 0)) or (plane == 'y' and exponents == (0, 1, 0, 0)): - - def transverse_dipolar_impedance_tesla_cavity(freq): - return 2 / ((2 * np.pi * freq / c_light) * a ** 2) * 1j * free_space_impedance / ( - np.pi * (2 * np.pi * freq / c_light) * a ** 2) * ( - 1 + (1 + 1j) * alpha * period_length / a * np.sqrt(np.pi / ((2 * np.pi * freq / c_light) * g))) ** (-1) - - def transverse_dipolar_wake_tesla_cavity(time): - return ((4 * free_space_impedance * c_light * s00) / (np.pi * a ** 4) * np.heaviside(time, 0) * - (1 - (1 + np.sqrt(c_light * time / s00)) * np.exp(-np.sqrt(c_light * time / s00)))) - - impedance = transverse_dipolar_impedance_tesla_cavity - wake = transverse_dipolar_wake_tesla_cavity - else: - print("Warning: tesla cavity impedance not implemented for component {}{}. Set to zero".format(plane, - exponents)) - - def zero_function(x): - return np.zeros_like(x) - - impedance = zero_function - wake = zero_function - - return Component(impedance=np.vectorize(lambda f: impedance(f)), - wake=np.vectorize(lambda t: wake(t)), - plane=plane, source_exponents=exponents[:2], test_exponents=exponents[2:]) - - -def _integral_stupakov(half_gap_small: float, half_gap_big: float, - half_width: float, g_index: int, g_power: int, - approximate_integrals: bool = False): - """ - Computes the Stupakov's integral for a rectangular linear taper. - Note that (g')^2 has been taken out of the integral. See Phys. Rev. STAB 10, 094401 (2007) - :param half_gap_small: the small half-gap of the taper - :param half_gap_big: the large half-gap of the taper - :param half_width: the half-width of the taper - :param g_index: the indice of the G function to use (0 is for G0=F in Stupakov's paper) - :param g_power: is the power to which 1/g is taken - :param approximate_integrals: use approximated formulas to compute the integrals. - It can be used if one assumes small half_gap_big/half_width ratio << 1 - """ - def _integrand_stupakov(g): - """ - Computes the integrand for the Stupakov integral - :param g: the half-gap of the taper - """ - x = g/half_width # half-gap over half-width ratio - - def g_addend(m): - """ - Computes the m-th addend of series defining F (when g_index=0) or G_[g_index] function from Stupakov's formulas for a - rectangular linear taper at a given x=half-gap/half-width. m can be an array, in which case the function returns an array. - See Phys. Rev. STAB 10, 094401 (2007) - :param m: the index of the addend (it can be an array) - """ - - if g_index == 0: - val = (2 * m + 1) * np.pi * x / 2 - res = ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * ( - 2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2) / (2 * m + 1) - - return res - - elif g_index == 1: - val = (2 * m + 1) * np.pi * x / 2 - res = x ** 3 * (2 * m + 1) * (4 * np.exp(-2 * val)) * (1 + np.exp(-2 * val)) / ((1 - np.exp(-2 * val)) ** 3) - return res - - elif g_index == 2: - val = (2 * m + 1) * np.pi * x / 2 - res = x ** 2 * (2 * m + 1) * ( - ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * (2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2)) - return res - - elif g_index == 3: - val = m * np.pi * x - res = x ** 2 * (2 * m) * ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * ( - 2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2) - return res - - else: - raise ValueError("Pb in g_addend for Stupakov's formulas: g_index not 0, 1, 2 or 3 !") - - # Computes F (when g_index=0) or G_[g_index] function from Stupakov's formulas. - g_stupakov = 0. - old_g_stupakov = 1. - eps = 1e-5 # relative precision of the summation - incr = 10 # increment for sum computation - m = np.arange(incr) - m_limit = 1e6 - - while (abs(g_stupakov - old_g_stupakov) > eps * abs(g_stupakov)) and (m[-1] < m_limit): - g_array = g_addend(m) - old_g_stupakov = g_stupakov - g_stupakov += np.sum(g_array) - m += incr - - if m[-1] >= m_limit: - print("Warning: maximum number of elements reached in g_stupakov!", m[-1], x, ", err=", - abs((g_stupakov - old_g_stupakov) / g_stupakov)) - - return g_stupakov / (g ** g_power) - - if approximate_integrals: - if g_index == 0 and g_power == 0: - i = 7. * zeta(3, 1) / (2. * np.pi ** 2) * ( - half_gap_big - half_gap_small) # (zeta(3.,1.) is Riemann zeta function at x=3) - - elif g_index == 1 and g_power == 3: - i = (1. / (half_gap_small ** 2) - 1. / (half_gap_big ** 2)) / ( - 2. * np.pi) # approx. integral - - elif g_index in [2,3] and g_power == 2: - i = (1. / half_gap_small - 1. / half_gap_big) / (np.pi ** 2) - - else: - raise ValueError("Wrong values of g_index and g_power in _integral_stupakov") - - else: - # computes numerically the integral instead of using its approximation - i, err = integrate.quadrature(_integrand_stupakov, half_gap_small, half_gap_big, - tol=1.e-3, maxiter=200, vec_func=False) - - return i - - -def shunt_impedance_flat_taper_stupakov_formula(half_gap_small: float, half_gap_big: float, taper_slope: float, - half_width: float, cutoff_frequency: float = None, - component_id: str = None, approximate_integrals: bool = False) -> float: - """ - Computes the shunt impedance in Ohm(/m if not longitudinal) of one single rectangular linear taper using Stupakov's - formulae (Phys. Rev. STAB 10, 094401 - 2007), multiplied by Z0*c/(4*pi) to convert to SI units. - Taper is in the vertical plane (i.e. the change of half-gap is in vertical, but the width is along the horizontal plane - - this is typical of a taper going towards a vertical collimator with horizontal jaws). - We use here half values for geometrical parameters (half-gap and half-width) whereas Stupakov's paper is expressed - with full values. This does not make any difference except for an additional factor 4 here for longitudinal impedance. - - The formula is valid under the conditions of low frequency and length of taper much larger than its transverse - dimensions. - - Note: one gets zero if component_id is not in [zlong, zxdip, zydip, zxqua, zyqua]. - - :param half_gap_small: small vertical half-gap - :param half_gap_big: large vertical half-gap - :param taper_slope: the slope of the taper - :param half_width: half width of the taper (constant) - :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) - :param component_id: a component name for which one computes the R/Q - (ex: zlong, zydip, zxqua, etc.) - :param approximate_integrals: use approximated formulas to compute the integrals. It can be used if one assumes - small half_gap_big/half_width ratio (<<1) - :return: the shunt impedance of the component. - """ - z_0 = mu_0 * c_light # free space impedance - - if cutoff_frequency is None and component_id == 'zlong': - raise ValueError("cutoff_frequency must be specified when component_id is 'zlong'") - - if component_id == 'zlong': - g_index = 0 - g_power = 0 - cst = 4. * mu_0 * cutoff_frequency / 2. # factor 4 due to use of half-gaps here - - elif component_id == 'zydip': - g_index = 1 - g_power = 3 - cst = z_0 * half_width * np.pi / 4. - - elif component_id == 'zxqua': - g_index = 2 - g_power = 2 - cst = -z_0 * np.pi / 4. - - elif component_id == 'zxdip': - g_index = 3 - g_power = 2 - cst = z_0 * np.pi / 4. - - elif component_id == 'zyqua': - g_index = 2 - g_power = 2 - cst = z_0 * np.pi / 4. - else: - # mock-up values - g_power = 0 - g_index = 0 - cst = 0 - - # computes the integral - i = _integral_stupakov(half_gap_small, half_gap_big, half_width, g_index, g_power, - approximate_integrals = approximate_integrals) - i *= taper_slope # put back g' factor that was dropped - - return cst * i - - -def create_flat_taper_stupakov_formula_component(half_gap_small: float, half_gap_big: float, taper_slope: float, - half_width: float, plane: str, exponents: Tuple[int, int, int, int], - cutoff_frequency: float = None) -> Component: - """ - Creates a component using the flat taper Stupakov formula - :param half_gap_small: small vertical half-gap - :param half_gap_big: large vertical half-gap - :param taper_slope: the slope of the taper - :param half_width: half width of the taper (constant) - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) - :return: A component object of a flat taper - """ - component_id = get_component_name(True, plane, exponents) - - r_shunt = shunt_impedance_flat_taper_stupakov_formula(half_gap_small=half_gap_small, half_gap_big=half_gap_big, - taper_slope=taper_slope, half_width=half_width, - cutoff_frequency=cutoff_frequency, component_id=component_id) - - return create_resonator_component(plane=plane, exponents=exponents, r=r_shunt, q=1, f_r=cutoff_frequency) - - -def create_flat_taper_stupakov_formula_element(half_gap_small: float, half_gap_big: float, taper_slope: float, - half_width: float, beta_x: float, beta_y: float, - cutoff_frequency: float = None, - component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), - name: str = "Flat taper", tag: str = "", - description: str = "") -> Element: - """ - Creates an element using the flat taper Stupakov formula - :param half_gap_small: small vertical half-gap - :param half_gap_big: large vertical half-gap - :param taper_slope: the slope of the taper - :param half_width: half width of the taper (constant) - :param beta_x: The beta function in the x-plane at the position of the taper - :param beta_y: The beta function in the y-plane at the position of the taper - :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) - :param component_ids: a list of components to be computed - :param name: A user-specified name for the Element - :param tag: A string to tag the Element - :param description: A description for the Element - :return: An Element object representing the flat taper - """ - length = (half_gap_big - half_gap_small) / taper_slope - - components = [] - for component_id in component_ids: - _, plane, exponents = component_names[component_id] - components.append(create_flat_taper_stupakov_formula_component(half_gap_small=half_gap_small, - half_gap_big=half_gap_big, - taper_slope=taper_slope, - half_width=half_width, - plane=plane, exponents=exponents, - cutoff_frequency=cutoff_frequency, - )) - - return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, - description=description) +from xwakes.wit.devices import * diff --git a/pywit/element.py b/pywit/element.py index cd1582ca..e9a6c3d9 100644 --- a/pywit/element.py +++ b/pywit/element.py @@ -1,310 +1 @@ -from __future__ import annotations - -from pywit.component import Component, Union - -from typing import List -from collections import defaultdict - -from scipy.special import comb -import numpy as np -import copy - -class Element: - def __init__(self, length: float = 0, beta_x: float = 0, beta_y: float = 0, - components: List[Component] = None, name: str = "Unnamed Element", - tag: str = "", description: str = ""): - """ - The initialization function of the Element class. - :param length: The length of the element, must be specified for valid initialization - :param beta_x: The size of the beta function in the x-plane at the position of the Element, must be specified - for valid initialization - :param beta_y: The size of the beta function in the y-plane at the position of the Element, must be specified - for valid initialization - :param components: A list of the Components corresponding to the Element being initialized. If the list contains - multiple Components with the same values for (plane, source_exponents, test_exponents), and error is thrown. If - the list is not specified, the Element is initialized with an empty components-list. - :param name: A user-specified name of the Element - :param tag: A string corresponding to a specific component - """ - assert length > 0, "The element must have a specified valid length" - assert beta_x > 0 and beta_y > 0, "The element must have valid specified beta_x and beta_y values" - self.length = length - self.beta_x = beta_x - self.beta_y = beta_y - self.name = name - self.tag = tag - self.description = description - if components: - comp_dict = defaultdict(int) - for c in components: - comp_dict[(c.plane, c.source_exponents, c.test_exponents)] += c - components = comp_dict.values() - self.components = sorted(components, key=lambda x: (x.plane, x.source_exponents, x.test_exponents)) - else: - self.components: List[Component] = [] - - def rotated(self, theta: float, rotate_beta: bool = False) -> Element: - """ - Returns a copy if the self-Element which has been rotated counterclockwise in the transverse plane by an - angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y values of the new Element - are rotated correspondingly. - :param theta: The angle, in radians, the Element is to be rotated by - :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new Element should also - be rotated by theta. Enabled by default. - :return: A newly initialized copy of the self-element which has been rotated counterclockwise in the - transverse plane by an angle theta. - """ - # Confines theta to the interval [0, 2 * pi) - theta %= 2 * np.pi - # Precalculates the sine and cosine of theta to save time - costheta, sintheta = np.cos(theta), np.sin(theta) - # Creates a dictionary where the keys are Component attributes and the values are Components with those - # corresponding attributes. That way, Components created by the rotation can always be added to other compatible - # Components. - rotated_components = defaultdict(int) - - # Iterates through all the components of the self-Element - for cmp in self.components: - # Defines the x- and y-coefficients depending on what plane the Component is in - if cmp.plane == 'x': - coefx, coefy = costheta, sintheta - else: - coefx, coefy = -sintheta, costheta - # For compactness, a, b, c and d are used to refer to the source_exponents and the test_exponents - # of the component - a, b, c, d = cmp.source_exponents + cmp.test_exponents - for i in range(a + 1): - for j in range(b + 1): - for k in range(c + 1): - for l in range(d + 1): - # Calculates new a, b, c and d values for the new Component - new_a, new_b, new_c, new_d = i + j, a - i + b - j, k + l, c - k + d - l - # Product of binomial coefficients - binprod = int(comb(a, i, exact=True) * comb(b, j, exact=True) * - comb(c, k, exact=True) * comb(d, l, exact=True)) - # Multiply by power of cos and sin - coef = ((-1) ** (j + l)) * binprod * (costheta ** (i + b - j + k + d - l)) * \ - (sintheta ** (a - i + j + c - k + l)) - # Depending on if the component is in the longitudinal or transverse plane, one or two - # components are created and added to the element at the correct key in the - # rotated_components dictionary - if cmp.plane == 'z': - # If the component is scaled by less than 10 ** -6 we assume that it is zero - if abs(coef) > 1e-6: - rotated_components['z', new_a, new_b, new_c, new_d] += \ - coef * Component(impedance=cmp.impedance, wake=cmp.wake, plane='z', - source_exponents=(new_a, new_b), - test_exponents=(new_c, new_d)) - else: - if abs(coefx * coef) > 1e-6: - rotated_components['x', new_a, new_b, new_c, new_d] += \ - (coefx * coef) * Component(impedance=cmp.impedance, wake=cmp.wake, plane='x', - source_exponents=(new_a, new_b), - test_exponents=(new_c, new_d)) - if abs(coefy * coef) > 1e-6: - rotated_components['y', new_a, new_b, new_c, new_d] += \ - (coefy * coef) * Component(impedance=cmp.impedance, wake=cmp.wake, plane='y', - source_exponents=(new_a, new_b), - test_exponents=(new_c, new_d)) - - # New beta_x and beta_y values are defined if the rotate_beta flag is active - new_beta_x = ((costheta * np.sqrt(self.beta_x) - - sintheta * np.sqrt(self.beta_y)) ** 2) if rotate_beta else self.beta_x - new_beta_y = ((sintheta * np.sqrt(self.beta_x) + - costheta * np.sqrt(self.beta_y)) ** 2) if rotate_beta else self.beta_y - - # Initializes and returns a new element with parameters calculated above. Its Component list is a list of - # the values in the rotated_components dictionary, sorted by the key which is used consistently for - # Component comparisons - return Element(self.length, new_beta_x, new_beta_y, - sorted(list(rotated_components.values()), - key=lambda x: (x.plane, x.source_exponents, x.test_exponents)), - self.name, self.tag, self.description) - - def is_compatible(self, other: Element, verbose: bool = False) -> bool: - """ - Compares all non-components parameters of the two Elements and returns False if they are not all equal within - some tolerance. The Component lists of the two Elements also need to be of equal length in order for the - function to return True. - :param other: An element for the self-Element to be compared against - :param verbose: A flag which can be activated to give feedback on the values of the parameters which caused - the function to return False - :return: True if all non-components parameters are equal within some tolerance and the Components lists of the - two Elements are of the same length. False otherwise. - """ - if verbose: - if abs(self.length - other.length) > 1e-8: - print(f"Different lengths: {self.length} != {other.length}") - return False - if abs(self.beta_x - other.beta_x) > 1e-8: - print(f"Different beta_x: {self.beta_x} != {other.beta_x}") - return False - if abs(self.beta_y - other.beta_y) > 1e-8: - print(f"Different beta_y: {self.beta_y} != {other.beta_y}") - return False - if len(self.components) != len(other.components): - print(f"Different number of components: {len(self.components)} != {len(other.components)}") - return False - return True - else: - return all([abs(self.length - other.length) < 1e-8, - abs(self.beta_x - other.beta_x) < 1e-8, - abs(self.beta_y - other.beta_y) < 1e-8, - len(self.components) == len(other.components)]) - - def changed_betas(self, new_beta_x: float, new_beta_y: float) -> Element: - element_copy = copy.deepcopy(self) - x_ratio = self.beta_x / new_beta_x - y_ratio = self.beta_y / new_beta_y - element_copy.components = [((x_ratio ** c.power_x) * (y_ratio ** c.power_y)) * c for c in self.components] - element_copy.beta_x = new_beta_x - element_copy.beta_y = new_beta_y - return element_copy - - def __add__(self, other: Element) -> Element: - """ - Defines the addition operator for two objects of the class Element - :param self: The left addend - :param other: The right addend - :return: A new object of the class Element which represents the sum of the two addends. Its length is simply - the sum of the length of the addends, its beta functions are a weighted, by length, sum of the beta functions - of the addends, and its list of components include all of the components of the addends, added together where - possible. - """ - # Defines new attributes based on attributes of the addends - new_length = self.length + other.length - new_beta_x = (self.length * self.beta_x + other.length * other.beta_x) / new_length - new_beta_y = (self.length * self.beta_y + other.length * other.beta_y) / new_length - new_components = [] - - # Pre-calculates some ratios which will be used in multiple calculations - ratios = (self.beta_x / new_beta_x, self.beta_y / new_beta_y, - other.beta_x / new_beta_x, other.beta_y / new_beta_y) - - # i and j represent indices in the component lists of the left and right element respectively - i, j = 0, 0 - # This while loop iterates through pairs of components in the component-lists of the elements until - # one of the lists is exhausted. - while i < len(self.components) and j < len(other.components): - # If the two components are compatible for addition, they are weighted according to the beta functions - if self.components[i].is_compatible(other.components[j]): - comp1 = self.components[i] - power_x = comp1.power_x - power_y = comp1.power_y - left_coefficient = (ratios[0] ** power_x) * (ratios[1] ** power_y) - right_coefficient = (ratios[2] ** power_x) * (ratios[3] ** power_y) - new_components.append(left_coefficient * comp1 + right_coefficient * other.components[j]) - i += 1 - j += 1 - # If the left component is "less than" the right component, by some arbitrary pre-defined measure, we know - # that there cannot be a component in the right element compatible with our left component. Therefore, we - # simply add the left component to the list of new components - elif self.components[i] < other.components[j]: - left_coefficient = (ratios[0] ** self.components[i].power_x) * (ratios[1] ** self.components[i].power_y) - new_components.append(left_coefficient * self.components[i]) - i += 1 - else: - right_coefficient = ((ratios[2] ** other.components[j].power_x) * - (ratios[3] ** other.components[j].power_y)) - new_components.append(right_coefficient * other.components[j]) - j += 1 - - # When the while-loop above exits, there could still be unprocessed components remaining in either the - # component list of either the left- or the right element, but not both. We simply append any remaining - # components to our new_components - if i != len(self.components): - for c in self.components[i:]: - left_coefficient = (ratios[0] ** c.power_x) * (ratios[1] ** c.power_y) - new_components.append(left_coefficient * c) - elif j != len(other.components): - for c in other.components[j:]: - right_coefficient = (ratios[2] ** c.power_x) * (ratios[3] ** c.power_y) - new_components.append(right_coefficient * c) - - # Creates and returns a new element which represents the sum of the two added elements. - return Element(new_length, new_beta_x, new_beta_y, new_components) - - def __radd__(self, other: Union[int, Element]) -> Element: - """ - Implements the __rad__ method for the Element class. This is only done to facilitate the syntactically - practical use of the sum() method for Elements. sum(iterable) works by adding all of the elements of the - iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case - that the left addend of any addition involving an Element is not itself an Element, the resulting sum - is simply defined to be the right addend. - :param other: The left addend of an addition - :return: The sum of self and other if other is an Element, otherwise just self. - """ - # Checks if the left addend, other, is not an Element - if not isinstance(other, Element): - # In which case, the right addend is simply returned - return self - - # Otherwise, their sum is returned (by invocation of Component.__add__(self, other)) - return self + other - - def __mul__(self, scalar: float) -> Element: - """ - Implements the __mul__ method for the Element class. Defines the behavior of multiplication of an Element by - some scalar. - :param scalar: A scalar value to be multiplied with some Element (cannot be complex) - :return: A newly initialized Element which has the same beta_x and beta_y values as the self-Element, but has - its length and all of the Components in its Component list multiplied by the scalar - """ - # Multiplies the length of the self-element by the scalar - new_length = self.length * scalar - # Creates a new list of Components which is a copy of the Component list of the self-element except every - # Component is multiplied by the scalar - new_components = [c * scalar for c in self.components] - # Initializes and returns a new Element with the arguments defined above - return Element(new_length, self.beta_x, self.beta_y, new_components, self.name, self.tag, self.description) - - def __rmul__(self, scalar: float) -> Element: - """ - Generalizes scalar multiplication of Element to be possibly from left and right. Both of these operations - are identical. - :param scalar: A scalar value to be multiplied with some Element - :return: The result of calling Element.__mul__(self, scalar): A newly initialized Element which has the same - beta_x and beta_y values as the self-Element, but has its length and all of the Components in its Component list - multiplied by the scalar. - """ - # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function - return self * scalar - - def __eq__(self, other: Element) -> bool: - """ - Implements the __eq__ method for the Element class. Two Elements are designated as "equal" if the following - two conditions hold: - 1. They have equal attributes within some tolerance and the length of the components-lists need to be identical. - 2. For every pair of components in the components-lists of the two elements, the two Components need to - evaluate as equal. - This somewhat approximated empirical approach to the equality comparator aims to compensate for small - numerical/precision errors accumulated for two Elements which have taken different "paths" to what should - analytically be identical Elements. - Note that this comparator requires the evaluation of every pair of wake- and impedance functions in some - number of points for every component in each of the elements. As such, it can be, though not extremely, - somewhat computationally intensive, and should not be used excessively. - :param other: The right hand side of the equality comparator - :return: True if the two Elements have sufficiently close attributes and every pair of components evaluate as - equal by the __eq__ method for the Component class. - """ - # Verifies that the two elements have sufficiently close attributes and components-lists of the same length - if not self.is_compatible(other): - return False - - # Returns true if every pair of components in the two lists evaluate as true by the __eq__ method defined in - # the Component class - return all(c1 == c2 for c1, c2 in zip(self.components, other.components)) - - def __str__(self): - return f"{self.name} with parameters:\n" \ - f"Length:\t\t{self.length}\n" \ - f"Beta_x:\t\t{self.beta_x}\n" \ - f"Beta_y:\t\t{self.beta_y}\n" \ - f"#components:\t{len(self.components)}" - - def get_component(self, type_string: str): - for comp in self.components: - if comp.get_shorthand_type() == type_string: - return comp - - raise KeyError(f"'{self.name}' has no component of the type '{type_string}'.") +from xwakes.wit.element import * diff --git a/pywit/elements_group.py b/pywit/elements_group.py index 87a7275a..1a16666d 100644 --- a/pywit/elements_group.py +++ b/pywit/elements_group.py @@ -1,178 +1 @@ -from __future__ import annotations - -import numpy as np - -from pywit.element import Element -from pywit.component import Component, Union - -from typing import List - - -class ElementsGroup(Element): - """ - A class used to store many elements of the same kind (i.e. Collimators, Roman Pots, Broadband resonators...). - Each of these groups require different handling of the input files and impedance computations (for some we use IW2D, - while for others we simply read the wake from a file), therefore this should be used as a base class from which - specific classes are derived. - """ - def __init__(self, elements_list: List[Element], name: str = "Unnamed Element", tag: str = "", - description: str = ""): - """ - The initialization function of the ElementsGroup class. - :param elements_list: The list of elements in this group - :param name: A user-specified name of the group - :param tag: An optional keyword that can be used to help grouping different ElementGroup's - :param description: A user-specified description of the group - """ - if len(elements_list) == 0: - raise ValueError('Elements_list cannot be empty') - - self.elements_list = elements_list - - sum_element = sum(elements_list) - - # Initialize the components by summing all components with the same values of (plane, source_exponents, - # test_exponents) - super().__init__(components=sum_element.components, name=name, tag=tag, description=description, - length=sum_element.length, beta_x=sum_element.beta_x, beta_y=sum_element.beta_y) - - def __add__(self, other: Union[Element, ElementsGroup]) -> ElementsGroup: - """ - Defines the addition operator for two objects of the class ElementsGroup or for an ElementsGroup and an Element. - :param self: The left addend - :param other: The right addend - :return: A new object of the class Element which represents the sum of the two addends. For two ElementsGroups - it is an Element given by the sum of all the elements in the two groups, while for an ElementsGroup and an - Element it is the Element given by the sum of the Element and all the elements in the ElementsGroup. - """ - if isinstance(other, Element): - elements_list = self.elements_list + [other] - elif isinstance(other, ElementsGroup): - elements_list = self.elements_list + other.elements_list - else: - raise TypeError("An ElementsGroup can only be summed with an ElementGroup or an Element") - - return ElementsGroup(elements_list, name=self.name, tag=self.tag, description=self.description) - - def __radd__(self, other: Union[int, Element, ElementsGroup]) -> ElementsGroup: - """ - Implements the __rad__ method for the ElementsGroup class. This is only done to facilitate the syntactically - practical use of the sum() method for ElementsGroup. sum(iterable) works by adding all of the elements of the - iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case - that the left addend of any addition involving an Element is not itself an Element, the resulting sum - is simply defined to be the right addend. - :param other: The left addend of an addition - :return: The sum of self and other if other is an ElementGroup or an Element, otherwise just self. - """ - if type(other) == int and other != 0: - raise ValueError("ElementsGroup right addition can only be performed with a zero addend and it is only " - "implemented to enable the sum on a list of ElementsGroup's") - # Checks if the left addend, other, is not an Element - if not (isinstance(other, ElementsGroup) or isinstance(other, Element)): - # In which case, the right addend is simply returned - return self - - # Otherwise, their sum is returned (by invocation of ElementsGroup.__add__(self, other)) - return self + other - - def __mul__(self, scalar: float) -> ElementsGroup: - """ - Implements the __mul__ method for the ElementsGroup class. Defines the behavior of multiplication of an Element - by some scalar. - :param scalar: A scalar value to be multiplied with some Element (cannot be complex) - :return: A newly initialized GroupElement in which every element is multiplied by the scalar. - """ - mult_elements = [] - for element in self.elements_list: - mult_elements.append(element * scalar) - - return ElementsGroup(mult_elements, name=self.name, tag=self.tag, description=self.description) - - def __rmul__(self, scalar: float) -> ElementsGroup: - """ - Generalizes scalar multiplication of ElementsGroup to be possibly from left and right. Both of these operations - are identical. - :param scalar: A scalar value to be multiplied with some ElementsGroup - :return: The result of calling ElementsGroup.__mul__(self, scalar): A newly initialized GroupElement in which - every element is multiplied by the scalar. - """ - # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function - return self * scalar - - def __eq__(self, other: ElementsGroup) -> bool: - """ - Implements the __eq__ method for the Element class. Two Elements are designated as "equal" if corresponding - elements in the elements lists are equal and if all the other non-components parameters are equal up to a small - tolerance - """ - if len(self.elements_list) != len(other.elements_list): - return False - - # Verifies that the two elements have sufficiently close attributes and components-lists of the same length - if not self.is_compatible(other): - return False - - return all(e1 == e2 for e1, e2 in zip(self.elements_list, other.elements_list)) - - def __str__(self): - string = f"{self.name} ElementsGroup composed of the following elements\n" - for element in self.elements_list: - string += str(element) - string += "============\n" - return string - - def rotated_element(self, name: str, theta: float, rotate_beta: bool = False) -> ElementsGroup: - """ - Returns a copy if the self-ElementsGroup in which a user-specified element has been rotated counterclockwise - in the transverse plane by an angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y - values of the element in the new ElementsGroup are rotated correspondingly. - :param name: the name of the Element to be rotated - :param theta: The angle, in radians, the ElementsGroup is to be rotated by - :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new ElementsGroup - should also be rotated by theta. Enabled by default. - :return: A newly initialized copy of the self-ElementsGroup which has been rotated counterclockwise in the - transverse plane by an angle theta. - """ - rotated_elements = self.elements_list.copy() - found = False - for i, element in enumerate(rotated_elements): - if element.name == name: - rotated_elements[i] = element.rotated(theta, rotate_beta) - found = True - - assert found, "Element to rotate was not found in the group" - - return ElementsGroup(rotated_elements, name=self.name, tag=self.tag, description=self.description) - - def rotated(self, theta: float, rotate_beta: bool = False) -> ElementsGroup: - """ - Returns a copy if the self-ElementsGroup which has been rotated counterclockwise in the transverse plane by an - angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y values of the new ElementsGroup - are rotated correspondingly. - :param theta: The angle, in radians, the ElementsGroup is to be rotated by - :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new ElementsGroup - should also be rotated by theta. Enabled by default. - :return: A newly initialized copy of the self-ElementsGroup which has been rotated counterclockwise in the - transverse plane by an angle theta. - """ - rotated_elements = [] - - for element in rotated_elements: - rotated_elements.append(element.rotated(theta, rotate_beta)) - - return ElementsGroup(rotated_elements, name=self.name, tag=self.tag, description=self.description) - - def changed_betas(self, new_beta_x: float, new_beta_y: float) -> ElementsGroup: - elements_list_new = [] - for i_elem, elem in enumerate(self.elements_list): - elements_list_new.append(elem.changed_betas(new_beta_x, new_beta_y)) - - return ElementsGroup(elements_list_new, name=self.name, tag=self.tag, description=self.description) - - def get_element(self, name_string: str): - for element in self.elements_list: - if element.name == name_string: - return element - - raise KeyError(f"'{self.name}' has no element named '{name_string}'.") - +from xwakes.wit.elements_group import * diff --git a/pywit/interface.py b/pywit/interface.py index d3c33091..ec0bb6ce 100644 --- a/pywit/interface.py +++ b/pywit/interface.py @@ -1,813 +1 @@ -import os - -from pywit.component import Component -from pywit.element import Element - -import subprocess -from typing import Tuple, List, Optional, Dict, Any, Union -from dataclasses import dataclass -from pathlib import Path -from hashlib import sha256 - -import numpy as np -from yaml import load, BaseLoader -from scipy.interpolate import interp1d - -# A dictionary mapping the datafile-prefixes (as used in IW2D) to (is_impedance, plane, (a, b, c, d)) -# Where is impedance is True if the component in question is an impedance component, and False if it is a -# wake component, and a, b, c and d are the source and test exponents of the component -component_names = {'wlong': (False, 'z', (0, 0, 0, 0)), - 'wxdip': (False, 'x', (1, 0, 0, 0)), - 'wydip': (False, 'y', (0, 1, 0, 0)), - 'wxqua': (False, 'x', (0, 0, 1, 0)), - 'wyqua': (False, 'y', (0, 0, 0, 1)), - 'wxcst': (False, 'x', (0, 0, 0, 0)), - 'wycst': (False, 'y', (0, 0, 0, 0)), - 'zlong': (True, 'z', (0, 0, 0, 0)), - 'zxdip': (True, 'x', (1, 0, 0, 0)), - 'zydip': (True, 'y', (0, 1, 0, 0)), - 'zxqua': (True, 'x', (0, 0, 1, 0)), - 'zyqua': (True, 'y', (0, 0, 0, 1)), - 'zxcst': (True, 'x', (0, 0, 0, 0)), - 'zycst': (True, 'y', (0, 0, 0, 0))} - -# The parent directory of this file -IW2D_SETTINGS_PATH = Path.home().joinpath('pywit').joinpath('config').joinpath('iw2d_settings.yaml') - - -def get_component_name(is_impedance, plane, exponents): - """ - Get the component name from is_impedance, plane and exponents (doing the - reverse operation of the dictionary in component_names) - :param is_impedance: True for impedance component, False for wake - :param plane: plane ('x', 'y' or 'z') - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :return: str with component name (e.g. 'zydip' or 'wxqua') - """ - comp_list = [comp_name for comp_name, v in component_names.items() - if v == (is_impedance, plane, exponents)] - if len(comp_list) != 1: - raise ValueError(f"({is_impedance},{plane},{exponents}) cannot be found in" - " the values of component_names dictionary") - - return comp_list[0] - - -def get_iw2d_config_value(key: str) -> Any: - with open(IW2D_SETTINGS_PATH, 'r') as file: - config = load(file, Loader=BaseLoader) - - return config[key] - - -def import_data_iw2d(directory: Union[str, Path], - common_string: str) -> List[Tuple[bool, str, Tuple[int, int, int, int], np.ndarray]]: - """ - Imports data on the format generated by the IW2D library and prepares it for construction of Components and - Elements in PyWIT - :param directory: The directory where the .dat files are located. All .dat files must be in the root of this - directory - :param common_string: A string preceding ".dat" in the filenames of all files to be imported - :return: A list of tuples, one for each imported file, on the form (is_impedance, plane, (a, b, c, d), data), - where data is a numpy array with 2 or 3 columns, one for each column of the imported datafile. - """ - # The tuples are iteratively appended to this array - component_recipes = [] - - # Keeps track of what combinations of (is_impedance, plane, exponents) have been imported to avoid duplicates - seen_configs = [] - - # A list of all of the filenames in the user-specified directory - filenames = os.listdir(directory) - for i, filename in enumerate(filenames): - # If the string preceding ".dat" in the filename does not match common_string, or if the first 5 letters - # of the filename are not recognized as a type of impedance/wake, the file is skipped - if filename[-4 - len(common_string):-4] != common_string or filename[:5].lower() not in component_names: - continue - - # The values of is_impedance, plane and exponents are deduced from the first 5 letters of the filename using - # the component_names-dictionary - is_impedance, plane, exponents = component_names[filename[:5].lower()] - - # Validates that the combination of (is_impedance, plane, exponents) is unique - assert (is_impedance, plane, exponents) not in seen_configs, \ - f"The {'impedance' if is_impedance else 'wake'} files " \ - f"'{filename}' and '{filenames[seen_configs.index((is_impedance, plane, exponents))]}' " \ - f"both correspond to the {plane}-plane with exponents {exponents}." - seen_configs.append((is_impedance, plane, exponents)) - - # Loads the data from the file as a numpy array - data = np.loadtxt(f"{directory}/{filename}", delimiter=" ", skiprows=1) - - # Appends the constructed tuple to component_recipes - component_recipes.append((is_impedance, plane, exponents, data)) - - # Validates that at least one file in the directory matched the user-specified common_string - assert component_recipes, f"No files in '{directory}' matched the common string '{common_string}'." - return component_recipes - - -def create_component_from_data(is_impedance: bool, plane: str, exponents: Tuple[int, int, int, int], - data: np.ndarray, relativistic_gamma: float, - extrapolate_to_zero: bool = False) -> Component: - """ - Creates a Component from a component recipe, e.g. as generated by import_data_iw2d - :param is_impedance: a bool which is True if the component to be generated is an impedance component, and False - if it is a wake component - :param plane: the plane of the component - :param exponents: the exponents of the component on the form (a, b, c, d) - :param data: a numpy-array with 2 or 3 columns corresponding to (frequency, Re[impedance], Im[impedance]) or - (position, Re[wake], Im[wake]), where the imaginary column is optional - :param relativistic_gamma: The relativistic gamma used in the computation of the data files. Necessary for - converting the position-data of IW2D into time-data for PyWIT - :param extrapolate_to_zero: a flag specifying if the data should be extrapolated to zero. If it is False (default - value), the data are extrapolated using the first and last value of the data. - :return: A Component object as specified by the input - """ - # Extracts the position/frequency column of the data array - x = data[:, 0] - - if not is_impedance: - # Converts position-data to time-data using Lorentz factor - x /= 299792458 * np.sqrt(1 - (1 / relativistic_gamma ** 2)) - - # Extracts the wake/values from the data array - y = data[:, 1] + (1j * data[:, 2] if data.shape[1] == 3 else 0) - - # Creates a callable impedance/wake function from the data array - if extrapolate_to_zero: - func = interp1d(x=x, y=y, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0, 0)) - else: - func = interp1d(x=x, y=y, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(y[0], y[-1])) - - # Initializes and returns a component based on the parameters provided - return Component(impedance=(func if is_impedance else None), - wake=(None if is_impedance else func), - plane=plane, - source_exponents=exponents[:2], - test_exponents=exponents[2:], ) - - -@dataclass(frozen=True, eq=True) -class Layer: - # The distance in mm of the inner surface of the layer from the reference orbit - thickness: float - dc_resistivity: float - resistivity_relaxation_time: float - re_dielectric_constant: float - magnetic_susceptibility: float - permeability_relaxation_frequency: float - - -@dataclass(frozen=True, eq=True) -class Sampling: - start: float - stop: float - # 0 = logarithmic, 1 = linear, 2 = both - scan_type: int - added: Tuple[float] - sampling_exponent: Optional[float] = None - points_per_decade: Optional[float] = None - min_refine: Optional[float] = None - max_refine: Optional[float] = None - n_refine: Optional[float] = None - - -# Define several dataclasses for IW2D input elements. We must split mandatory -# and optional arguments into private dataclasses to respect the resolution -# order. The public classes RoundIW2DInput and FlatIW2D input inherit from -# from the private classes. -# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses - -@dataclass(frozen=True, eq=True) -class _IW2DInputBase: - machine: str - length: float - relativistic_gamma: float - calculate_wake: bool - f_params: Sampling - - -@dataclass(frozen=True, eq=True) -class _IW2DInputOptional: - z_params: Optional[Sampling] = None - long_factor: Optional[float] = None - wake_tol: Optional[float] = None - freq_lin_bisect: Optional[float] = None - comment: Optional[str] = None - - -@dataclass(frozen=True, eq=True) -class IW2DInput(_IW2DInputOptional, _IW2DInputBase): - pass - - -@dataclass(frozen=True, eq=True) -class _RoundIW2DInputBase(_IW2DInputBase): - layers: Tuple[Layer] - inner_layer_radius: float - # (long, xdip, ydip, xquad, yquad) - yokoya_factors: Tuple[float, float, float, float, float] - - -@dataclass(frozen=True, eq=True) -class _RoundIW2DInputOptional(_IW2DInputOptional): - pass - - -@dataclass(frozen=True, eq=True) -class RoundIW2DInput(_RoundIW2DInputOptional, _RoundIW2DInputBase): - pass - - -@dataclass(frozen=True, eq=True) -class _FlatIW2DInputBase(_IW2DInputBase): - top_bottom_symmetry: bool - top_layers: Tuple[Layer] - top_half_gap: float - - -@dataclass(frozen=True, eq=True) -class _FlatIW2DInputOptional(_IW2DInputOptional): - bottom_layers: Optional[Tuple[Layer]] = None - bottom_half_gap: Optional[float] = None - - -@dataclass(frozen=True, eq=True) -class FlatIW2DInput(_FlatIW2DInputOptional, _FlatIW2DInputBase): - pass - - -def _iw2d_format_layer(layer: Layer, n: int) -> str: - """ - Formats the information describing a single layer into a string in accordance with IW2D standards. - Intended only as a helper-function for create_iw2d_input_file. - :param layer: A Layer object - :param n: The 1-indexed index of the layer - :return: A string on the correct format for IW2D - """ - return (f"Layer {n} DC resistivity (Ohm.m):\t{layer.dc_resistivity}\n" - f"Layer {n} relaxation time for resistivity (ps):\t{layer.resistivity_relaxation_time * 1e12}\n" - f"Layer {n} real part of dielectric constant:\t{layer.re_dielectric_constant}\n" - f"Layer {n} magnetic susceptibility:\t{layer.magnetic_susceptibility}\n" - f"Layer {n} relaxation frequency of permeability (MHz):\t{layer.permeability_relaxation_frequency / 1e6}\n" - f"Layer {n} thickness in mm:\t{layer.thickness * 1e3}\n") - - -def _iw2d_format_freq_params(params: Sampling) -> str: - """ - Formats the frequency-parameters of an IW2DInput object to a string in accordance with IW2D standards. - Intended only as a helper-function for create_iw2d_input_file. - :param params: Parameters specifying a frequency-sampling - :return: A string on the correct format for IW2D - """ - lines = [f"start frequency exponent (10^) in Hz:\t{np.log10(params.start)}", - f"stop frequency exponent (10^) in Hz:\t{np.log10(params.stop)}", - f"linear (1) or logarithmic (0) or both (2) frequency scan:\t{params.scan_type}"] - - if params.sampling_exponent is not None: - lines.append(f"sampling frequency exponent (10^) in Hz (for linear):\t{np.log10(params.sampling_exponent)}") - - if params.points_per_decade is not None: - lines.append(f"Number of points per decade (for log):\t{params.points_per_decade}") - - if params.min_refine is not None: - lines.append(f"when both, fmin of the refinement (in THz):\t{params.min_refine / 1e12}") - - if params.max_refine is not None: - lines.append(f"when both, fmax of the refinement (in THz):\t{params.max_refine / 1e12}") - - if params.n_refine is not None: - lines.append(f"when both, number of points in the refinement:\t{params.n_refine}") - - lines.append(f"added frequencies [Hz]:\t{' '.join(str(f) for f in params.added)}") - - return "\n".join(lines) + "\n" - - -def _iw2d_format_z_params(params: Sampling) -> str: - """ - Formats the position-parameters of an IW2DInput object to a string in accordance with IW2D standards. - Intended only as a helper-function for create_iw2d_input_file. - :param params: Parameters specifying a position-sampling - :return: A string on the correct format for IW2D - """ - lines = [f"linear (1) or logarithmic (0) or both (2) scan in z for the wake:\t{params.scan_type}"] - - if params.sampling_exponent is not None: - lines.append(f"sampling distance in m for the linear sampling:\t{params.sampling_exponent}") - - if params.min_refine is not None: - lines.append(f"zmin in m of the linear sampling:\t{params.min_refine}") - - if params.max_refine is not None: - lines.append(f"zmax in m of the linear sampling:\t{params.max_refine}") - - if params.points_per_decade is not None: - lines.append(f"Number of points per decade for the logarithmic sampling:\t{params.points_per_decade}") - - lines.append(f"exponent (10^) of zmin (in m) of the logarithmic sampling:\t{np.log10(params.start)}") - lines.append(f"exponent (10^) of zmax (in m) of the logarithmic sampling:\t{np.log10(params.stop)}") - lines.append(f"added z [m]:\t{' '.join(str(z) for z in params.added)}") - - return "\n".join(lines) + "\n" - - -def create_iw2d_input_file(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], filename: Union[str, Path]) -> None: - """ - Writes an IW2DInput object to the specified filename using the appropriate format for interfacing with the IW2D - software. - :param iw2d_input: An IW2DInput object to be written - :param filename: The filename (including path) of the file the IW2DInput object will be written to - :return: Nothing - """ - # Creates the input-file at the location specified by filename - file = open(filename, 'w') - - file.write(f"Machine:\t{iw2d_input.machine}\n" - f"Relativistic Gamma:\t{iw2d_input.relativistic_gamma}\n" - f"Impedance Length in m:\t{iw2d_input.length}\n") - - # Just pre-defining layers to avoid potentially unbound variable later on - layers = [] - if isinstance(iw2d_input, RoundIW2DInput): - file.write(f"Number of layers:\t{len(iw2d_input.layers)}\n" - f"Layer 1 inner radius in mm:\t{iw2d_input.inner_layer_radius * 1e3}\n") - layers = iw2d_input.layers - elif isinstance(iw2d_input, FlatIW2DInput): - if iw2d_input.bottom_layers: - print("WARNING: bottom layers of IW2D input object are being ignored because the top_bottom_symmetry flag " - "is enabled") - file.write(f"Number of upper layers in the chamber wall:\t{len(iw2d_input.top_layers)}\n") - if iw2d_input.top_layers: - file.write(f"Layer 1 inner half gap in mm:\t{iw2d_input.top_half_gap * 1e3}\n") - layers = iw2d_input.top_layers - - for i, layer in enumerate(layers): - file.write(_iw2d_format_layer(layer, i + 1)) - - if isinstance(iw2d_input, FlatIW2DInput) and not iw2d_input.top_bottom_symmetry: - file.write(f"Number of lower layers in the chamber wall:\t{len(iw2d_input.bottom_layers)}\n") - if iw2d_input.bottom_layers: - file.write(f"Layer -1 inner half gap in mm:\t{iw2d_input.bottom_half_gap * 1e3}\n") - for i, layer in enumerate(iw2d_input.bottom_layers): - file.write(_iw2d_format_layer(layer, -(i + 1))) - - if isinstance(iw2d_input, FlatIW2DInput): - file.write(f"Top bottom symmetry (yes or no):\t{'yes' if iw2d_input.top_bottom_symmetry else 'no'}\n") - - file.write(_iw2d_format_freq_params(iw2d_input.f_params)) - if iw2d_input.z_params is not None: - file.write(_iw2d_format_z_params(iw2d_input.z_params)) - - if isinstance(iw2d_input, RoundIW2DInput): - file.write(f"Yokoya factors long, xdip, ydip, xquad, yquad:\t" - f"{' '.join(str(n) for n in iw2d_input.yokoya_factors)}\n") - - for desc, val in zip(["factor weighting the longitudinal impedance error", - "tolerance (in wake units) to achieve", - "frequency above which the mesh bisecting is linear [Hz]", - "Comments for the output files names"], - [iw2d_input.long_factor, iw2d_input.wake_tol, iw2d_input.freq_lin_bisect, iw2d_input.comment]): - if val is not None: - file.write(f"{desc}:\t{val}\n") - - file.close() - - -def check_already_computed(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], - name: str) -> Tuple[bool, str, Union[str, Path]]: - """ - Checks if a simulation with inputs iw2d_input is already present in the hash database. - :param iw2d_input: an iw2d input object - :param name: the name of the object - :return already_computed: boolean indicating if the iw2d_inputs have been already - computed - :return input_hash: string with the hash key corresponding to the inputs - :return working_directory: the path to the directory where the files were put, which is built as - `// - """ - projects_path = Path(get_iw2d_config_value('project_directory')) - - # initialize read ready to all False for convenience - # create the hash key - input_hash = sha256(iw2d_input.__str__().encode()).hexdigest() - - # we have three levels of directories: the first two are given by the first and second letters of the hash keys, - # the third is given by the rest of the hash keys. - directory_level_1 = projects_path.joinpath(input_hash[0:2]) - directory_level_2 = directory_level_1.joinpath(input_hash[2:4]) - working_directory = directory_level_2.joinpath(input_hash[4:]) - - already_computed = True - - # check if the directories exist. If they do not exist we create - if not os.path.exists(directory_level_1): - already_computed = False - os.mkdir(directory_level_1) - - if not os.path.exists(directory_level_2): - already_computed = False - os.mkdir(directory_level_2) - - if not os.path.exists(working_directory): - already_computed = False - os.mkdir(working_directory) - - components = [] - if not iw2d_input.calculate_wake: - for component in component_names.keys(): - # the ycst component is only given in the case of a flat chamber and the x component is never given - if component.startswith('z') and 'cst' not in component: - components.append(component) - if isinstance(iw2d_input, FlatIW2DInput): - components.append('zycst') - else: - for component in component_names.keys(): - # if the wake is computed, all keys from component_names dict are added, except the constant impedance/wake - # in first instance. If the simulation is a flat chamber, we add the vertical constant impedance/wake - if 'cst' not in component: - components.append(component) - if isinstance(iw2d_input, FlatIW2DInput): - components.append('wycst') - components.append('zycst') - - # The simulation seems to have been already computed, but we check if all the components of the impedance - # wake have been computed. If not, the computation will be relaunched - if already_computed: - # this list also includes the input file but it doesn't matter - computed_components = [name[0:5].lower() for name in os.listdir(working_directory)] - - for component in components: - if component not in computed_components: - already_computed = False - break - - return already_computed, input_hash, working_directory - - -def check_valid_hash_chunk(hash_chunk: str, length: int): - """ - Checks that the hash_chunk string can be the part of an hash key of given length. This means that hash_chunk must - have the right length and it must be a hexadecimal string - :param hash_chunk: the string to be checked - :param length: the length which the hash chunk should have - :return: True if hash_chunk is valid, False otherwise - """ - if len(hash_chunk) != length: - return False - - # check if the hash is an hexadecimal string - try: - int(hash_chunk, 16) - return True - except ValueError: - return False - - -def check_valid_working_directory(working_directory: Path): - """ - Checks if working_directory is valid. To be valid working directory must be of the form - `/hash[0:2]/hash[2:4]/hash[4:]` - :param working_directory: the path to the directory to be checked - :return: True if the working_directory is valid, False otherwise - """ - projects_path = Path(get_iw2d_config_value('project_directory')) - - if working_directory.parent.parent.parent != projects_path: - raise ValueError(f"The working directory must be located inside {projects_path}") - - return (check_valid_hash_chunk(working_directory.parent.parent.name, 2) and - check_valid_hash_chunk(working_directory.parent.name, 2) and - check_valid_hash_chunk(working_directory.name, 60)) - - -def add_iw2d_input_to_database(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], input_hash: str, - working_directory: Union[str, Path]): - """ - Add the iw2d input to the repository containing the simulations - :param iw2d_input: the input object of the IW2D simulation - :param input_hash: the hash key corresponding to the input - :param working_directory: the directory where to put the iw2d input file - """ - if type(working_directory) == str: - working_directory = Path(working_directory) - - if not check_valid_working_directory(working_directory): - raise ValueError("working directory is not in the right format. The right format is " - "`/hash[0:2]/hash[2:4]/hash[4:]`") - - directory_level_1 = working_directory.parent.parent - directory_level_2 = working_directory.parent - - if not os.path.exists(directory_level_1): - os.mkdir(directory_level_1) - if not os.path.exists(directory_level_2): - os.mkdir(directory_level_2) - - working_directory = directory_level_2.joinpath(input_hash[4:]) - - if not os.path.exists(working_directory): - os.mkdir(working_directory) - - create_iw2d_input_file(iw2d_input, working_directory.joinpath(f"input.txt")) - - -def create_element_using_iw2d(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], name: str, beta_x: float, beta_y: float, - tag: str = 'IW2D', extrapolate_to_zero: bool = False) -> Element: - """ - Create and return an Element using IW2D object. - :param iw2d_input: the IW2DInput object - :param name: the name of the Element - :param beta_x: the beta function value in the x-plane at the position of the Element - :param beta_y: the beta function value in the x-plane at the position of the Element - :param tag: a tag string for the Element - :param extrapolate_to_zero: a flag specifying if the data should be extrapolated to zero. If it is False (default - value), the data are extrapolated using the first and last value of the data. - :return: The newly computed Element - """ - assert " " not in name, "Spaces are not allowed in element name" - - assert verify_iw2d_config_file(), "The binary and/or project directories specified in config/iw2d_settings.yaml " \ - "do not exist or do not contain the required files and directories." - - # when looking for this IW2DInput in the database, the comment and the machine name don't necessarily need to be - # the same as the in the old simulation so we ignore it for creating the hash - iw2d_input_dict = iw2d_input.__dict__ - comment = iw2d_input_dict['comment'] - machine = iw2d_input_dict['machine'] - iw2d_input_dict['comment'] = '' - iw2d_input_dict['machine'] = '' - - # the path to the folder containing the IW2D executables - bin_path = Path(get_iw2d_config_value('binary_directory')) - # the path to the folder containing the database of already computed elements - - # check if the element is already present in the database and create the hash key corresponding to the IW2D input - already_computed, input_hash, working_directory = check_already_computed(iw2d_input, name) - - if already_computed: - print(f"The computation of '{name}' has already been performed with the exact given parameters. " - f"These results will be used to generate the element.") - - # if an element with the same inputs is not found inside the database, perform the computations and add the results - # to the database - if not already_computed: - add_iw2d_input_to_database(iw2d_input, input_hash, working_directory) - bin_string = ("wake_" if iw2d_input.calculate_wake else "") + \ - ("round" if isinstance(iw2d_input, RoundIW2DInput) else "flat") + "chamber.x" - subprocess.run(f'{bin_path.joinpath(bin_string)} < input.txt', shell=True, cwd=working_directory) - - # When the wake is computed with IW2D, a second set of files is provided by IW2D. These correspond to a "converged" - # simulation with double the number of mesh points for the wake. They files have the _precise suffix to their name. - # If the wake is computed, we retrieve these file to create the pywit element. - common_string = "_precise" if iw2d_input.calculate_wake else '' - - component_recipes = import_data_iw2d(directory=working_directory, common_string=common_string) - - iw2d_input_dict['comment'] = comment - iw2d_input_dict['machine'] = machine - - return Element(length=iw2d_input.length, - beta_x=beta_x, beta_y=beta_y, - components=[create_component_from_data(*recipe, relativistic_gamma=iw2d_input.relativistic_gamma, - extrapolate_to_zero=extrapolate_to_zero) - for recipe in component_recipes], - name=name, tag=tag, description='A resistive wall element created using IW2D') - - -def verify_iw2d_config_file() -> bool: - bin_path = Path(get_iw2d_config_value('binary_directory')) - projects_path = Path(get_iw2d_config_value('project_directory')) - if not bin_path.exists() or not projects_path.exists(): - return False - - contents = os.listdir(bin_path) - for filename in ('flatchamber.x', 'roundchamber.x', 'wake_flatchamber.x', 'wake_roundchamber.x'): - if filename not in contents: - return False - - return True - - -def _typecast_sampling_dict(d: Dict[str, str]) -> Dict[str, Any]: - added = [float(f) for f in d['added'].split()] if 'added' in d else [] - added = tuple(added) - scan_type = int(d['scan_type']) - d.pop('added'), d.pop('scan_type') - - new_dict = {k: float(v) for k, v in d.items()} - new_dict['added'] = added - new_dict['scan_type'] = scan_type - return new_dict - - -def _create_iw2d_input_from_dict(d: Dict[str, Any]) -> Union[FlatIW2DInput, RoundIW2DInput]: - is_round = d['is_round'].lower() in ['true', 'yes', 'y', '1'] - d.pop('is_round') - layers, inner_layer_radius, yokoya_factors = list(), float(), tuple() - top_layers, top_half_gap, bottom_layers, bottom_half_gap = list(), float(), None, None - - if is_round: - inner_layer_radius = d['inner_layer_radius'] - if 'layers' in d: - layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['layers']] - layers = [Layer(**kwargs) for kwargs in layers_dicts] - d.pop('layers') - else: - if 'top_layers' in d: - top_layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['top_layers']] - top_layers = [Layer(**kwargs) for kwargs in top_layers_dicts] - top_half_gap = d['top_half_gap'] - d.pop('top_layers') - if d['top_bottom_symmetry'].lower() in ['true', 'yes', 'y', '1']: - bottom_layers = None - else: - bottom_layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['bottom_layers']] - bottom_layers = [Layer(**kwargs) for kwargs in bottom_layers_dicts] - bottom_half_gap = d['bottom_half_gap'] - d.pop('bottom_layers') - - if 'yokoya_factors' in d: - yokoya_factors = tuple(float(x) for x in d['yokoya_factors'].split()) - d.pop('yokoya_factors') - - f_params = Sampling(**_typecast_sampling_dict(d['f_params'])) - z_params = Sampling(**_typecast_sampling_dict(d['z_params'])) \ - if d['calculate_wake'].lower() in ['true', 'yes', 'y', '1'] else None - - d.pop('f_params') - d.pop('z_params', None) - - transformations = { - 'machine': str, - 'length': float, - 'relativistic_gamma': float, - 'calculate_wake': lambda x: x.lower() in ['true', 'yes', 'y', '1'], - 'long_factor': float, - 'wake_tol': float, - 'freq_lin_bisect': float, - 'comment': str - } - - new_dict = {k: transformations[k](d[k]) if k in d else None for k in transformations} - - if is_round: - return RoundIW2DInput( - f_params=f_params, - z_params=z_params, - layers=tuple(layers), - inner_layer_radius=inner_layer_radius, - yokoya_factors=yokoya_factors, - **new_dict - ) - else: - return FlatIW2DInput( - f_params=f_params, - z_params=z_params, - top_bottom_symmetry=d['top_bottom_symmetry'].lower() in ['true', 'yes', 'y', '1'], - top_layers=tuple(top_layers), - top_half_gap=top_half_gap, - bottom_layers=bottom_layers, - bottom_half_gap=bottom_half_gap, - **new_dict - ) - - -def create_iw2d_input_from_yaml(name: str) -> Union[FlatIW2DInput, RoundIW2DInput]: - """ - Create a IW2DInput object from one of the inputs specified in the `pywit/config/iw2d_inputs.yaml` database - :param name: the name of the input which is read from the yaml database - :return: the newly initialized IW2DInput object - """ - path = Path.home().joinpath('pywit').joinpath('config').joinpath('iw2d_inputs.yaml') - with open(path) as file: - inputs = load(file, Loader=BaseLoader) - d = inputs[name] - - return _create_iw2d_input_from_dict(d) - - -def create_multiple_elements_using_iw2d(iw2d_inputs: List[IW2DInput], names: List[str], - beta_xs: List[float], beta_ys: List[float]) -> List[Element]: - """ - Create and return a list of Element's using a list of IW2D objects. - :param iw2d_inputs: the list of IW2DInput objects - :param names: the list of names of the Element's - :param beta_xs: the list of beta function values in the x-plane at the position of each Element - :param beta_ys: the list of beta function values in the x-plane at the position of each Element - :return: the list of newly computed Element's - """ - assert len(iw2d_inputs) == len(names) == len(beta_xs) == len(beta_ys), "All input lists need to have the same" \ - "number of elements" - - for name in names: - assert " " not in name, "Spaces are not allowed in element name" - - assert verify_iw2d_config_file(), "The binary and/or project directories specified in config/iw2d_settings.yaml " \ - "do not exist or do not contain the required files and directories." - - - from joblib import Parallel, delayed - elements = Parallel(n_jobs=-1, prefer='threads')(delayed(create_element_using_iw2d)( - iw2d_inputs[i], - names[i], - beta_xs[i], - beta_ys[i] - ) for i in range(len(names))) - - return elements - - -def create_htcondor_input_file(iw2d_input: IW2DInput, name: str, directory: Union[str, Path]) -> None: - exec_string = "" - if iw2d_input.calculate_wake: - exec_string += "wake_" - exec_string += ("round" if isinstance(iw2d_input, RoundIW2DInput) else "flat") + "chamber.x" - - text = f"executable = {exec_string}\n" \ - f"input = {name}_input.txt\n" \ - f"ID = $(Cluster).$(Process)\n" \ - f"output = $(ID).out\n" \ - f"error = $(ID).err\n" \ - f"log = $(Cluster).log\n" \ - f"universe = vanilla\n" \ - f"initialdir = \n" \ - f"when_to_transfer_output = ON_EXIT\n" \ - f'+JobFlavour = "tomorrow"\n' \ - f'queue' - - with open(directory, 'w') as file: - file.write(text) - - -def _verify_iw2d_binary_directory(ignore_missing_files: bool = False) -> None: - bin_path = Path(get_iw2d_config_value('binary_directory')) - if not ignore_missing_files: - filenames = ('flatchamber.x', 'roundchamber.x', 'wake_flatchamber.x', 'wake_roundchamber.x') - assert all(filename in os.listdir(bin_path) for filename in filenames), \ - "In order to utilize IW2D with PyWIT, the four binary files 'flatchamber.x', 'roundchamber.x', " \ - f"'wake_flatchamber.x' and 'wake_roundchamber.x' (as generated by IW2D) must be placed in the directory " \ - f"'{bin_path}'." - - -def _read_cst_data(filename: Union[str, Path]) -> np.ndarray: - with open(filename, 'r') as f: - lines = f.readlines() - data = [] - for line in lines: - try: - data.append([float(e) for e in line.strip().split()]) - except ValueError: - pass - - return np.asarray(data) - - -def load_longitudinal_impedance_datafile(path: Union[str, Path]) -> Component: - data = _read_cst_data(path) - fs = data[:, 0] - zs = data[:, 1] + 1j * data[:, 2] - func = interp1d(x=fs, y=zs, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) - return Component(impedance=func, plane='z', source_exponents=(0, 0), test_exponents=(0, 0)) - - -def load_transverse_impedance_datafile(path: Union[str, Path]) -> Tuple[Component, Component, Component, Component]: - data = _read_cst_data(path) - fs = data[:, 0] - zs = [data[:, 2 * i + 1] + 1j * data[:, 2 * i + 2] for i in range(4)] - components = tuple() - for i, z in enumerate(zs): - exponents = [int(j == i) for j in range(4)] - func = interp1d(x=fs, y=z, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) - components += (Component(impedance=func, - plane='x' if i % 2 == 0 else 'y', - source_exponents=(exponents[0], exponents[1]), - test_exponents=(exponents[2], exponents[3])),) - - return components - - -def load_longitudinal_wake_datafile(path: Union[str, Path]) -> Component: - data = _read_cst_data(path) - ts = data[:, 0] - ws = data[:, 1] * 1e15 - func = interp1d(x=ts, y=ws, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) - return Component(wake=func, plane='z', source_exponents=(0, 0), test_exponents=(0, 0)) - - -def load_transverse_wake_datafile(path: Union[str, Path]) -> Tuple[Component, Component, Component, Component]: - data = _read_cst_data(path) - ts = data[:, 0] - ws = [data[:, i] * 1e15 for i in range(1, 5)] - components = tuple() - for i, w in enumerate(ws): - exponents = [int(j == i) for j in range(4)] - func = interp1d(x=ts, y=w, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) - components += (Component(wake=func, - plane='x' if i % 2 == 0 else 'y', - source_exponents=(exponents[0], exponents[1]), - test_exponents=(exponents[2], exponents[3])),) - - return components - +from xwakes.wit.interface import * diff --git a/pywit/landau_damping.py b/pywit/landau_damping.py index 97127573..39882ab7 100644 --- a/pywit/landau_damping.py +++ b/pywit/landau_damping.py @@ -1,155 +1 @@ -import numpy as np -from scipy.special import exp1 -from scipy.optimize import newton - -from typing import Sequence, Tuple - - -def dispersion_integral_2d(tune_shift: np.ndarray, b_direct: float, b_cross: float, - distribution: str = 'gaussian'): - """ - Compute the dispersion integral in 2D from q complex tune shift, given the detuning coefficients (multiplied by - sigma). This is the integral vs Jx and Jy of Jx*dphi/dJx/(Q-bx*Jx-bxy*Jy-i0) (with phi the distribution function) - The transverse distribution can be 'gaussian' or 'parabolic'. - Note: for stability diagrams, use -1/dispersion_integral, and usually the convention is to plot -Im[Q] vs Re[Q]. - Reference: Berg-Ruggiero: https://cds.cern.ch/record/318826?ln=it - :param tune_shift: the complex tune shift - :param b_direct: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in - the x plane or $\alpha_y \sigma_y$ if working in the y plane) - :param b_cross: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in - the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) - :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' - :return: the dispersion integral - """ - if np.imag(tune_shift) == 0: - tune_shift = tune_shift - 1e-15 * 1j - - c = b_cross / b_direct - q = -tune_shift / b_direct - - if distribution == 'gaussian': - - i1 = (1 - c - (q + c - c * q) * np.exp(q) * exp1(q) + c * np.exp(q / c) * exp1(q / c)) / ((1 - c) ** 2) - - if np.isnan(i1): - i1 = 1. / q - (c + 2) / q ** 2 # asymptotic form for large q (assuming c is of order 1) - - elif distribution == 'parabolic': - - xi = q / 5. - - if np.abs(xi) > 100.: - # asymptotic form for large q (assuming c is of order 1) - i1 = 1. / q - (c + 2) / q ** 2 - else: - i1 = (((c + xi) ** 3 * np.log((1 + xi) / (c + xi)) + - (-1 + c) * (c * (c + 2 * c * xi + (-1 + 2 * c) * xi ** 2) + - (-1 + c) * xi ** 2 * (3 * c + xi + 2 * c * xi) * np.log(xi / (1 + xi)))) / - ((-1 + c) ** 2 * c ** 2)) - i1 = -i1 * 4. / 5. - # this is the same as in Scott Berg-Ruggiero CERN SL-AP-96-71 (AP) - - else: - raise ValueError("distribution can only be 'gaussian' or 'parabolic'") - - i = -i1 / b_direct - - # additional minus sign because we want the integral with dphi/dJx (derivative of distribution) on the - # numerator, so -[the one of Berg-Ruggiero] - return -i - - -def find_detuning_coeffs_threshold(tune_shift: complex, q_s: float, b_direct_ref: float, b_cross_ref: float, - fraction_of_qs_allowed_on_positive_side: float = 0.05, - distribution: str = 'gaussian', tolerance=1e-10): - """ - Compute the detuning coefficients (multiplied by sigma) corresponding to stability diagram threshold for a complex - tune shift. - It keeps fixed the ratio between b_direct_ref and b_cross_ref. - :param tune_shift: the tune shift for which the octupole threshold is computed - :param q_s: the synchrotron tune - :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in - the x plane or $\alpha_y \sigma_y$ if working in the y plane) - :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in - the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) - :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' - :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the - stability diagram), one can consider positive tune shift up to this fraction of q_s (default=5%) - :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding - and for the final check that the roots are actually proper roots. - :return: the detuning coefficients corresponding to the stability diagram threshold if the corresponding mode is - unstable, 0 if the corresponding mode is stable or np.nan if the threshold cannot be found (failure of Newton's - algorithm). - """ - # evaluate azimuthal mode number - l_mode = int(np.ceil(np.real(tune_shift) / q_s)) - if (l_mode - np.real(tune_shift) / q_s) > 1 - fraction_of_qs_allowed_on_positive_side: - l_mode -= 1 - # take away the shift from azimuthal mode number - tune_shift -= q_s * l_mode - - b_ratio = b_cross_ref/b_direct_ref - if tune_shift.imag < 0.: - - # function to solve (distance in imag. part w.r.t stab. diagram, as a function of oct. current) - def f(b_direct): - b_direct_i = b_direct - b_cross_i = b_ratio * b_direct - stab = [-1. / dispersion_integral_2d(t_s, b_direct_i, b_cross_i, distribution=distribution) for e in (-1, 1) - for t_s in b_direct_i * e * 10. ** np.arange(-3, 2, 0.01)[::e]] - # note: one has to reverse the table to get the interpolation right, for negative polarity (np.interp always - # wants monotonically increasing abscissae) - return tune_shift.imag - np.interp(tune_shift.real, np.real(stab)[::int(np.sign(b_direct_ref))], - np.imag(stab)[::int(np.sign(b_direct_ref))]) - - # Newton root finding - try: - b_direct_new = newton(f, b_direct_ref, tol=tolerance) - except RuntimeError: - b_direct_new = np.nan - else: - if np.abs(f(b_direct_new)) > tolerance: - b_direct_new = np.nan - else: - b_direct_new = 0. - - return b_direct_new, b_ratio*b_direct_new - - -def abs_first_item_or_nan(tup: Tuple): - if tup is not np.nan: - return abs(tup[0]) - else: - return np.nan - - -def find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts: Sequence[complex], q_s: float, b_direct_ref: float, - b_cross_ref: float, distribution: str = 'gaussian', - fraction_of_qs_allowed_on_positive_side: float = 0.05, - tolerance=1e-10): - """ - Compute the detuning coefficients corresponding to the most stringent stability diagram threshold for a sequence of - complex tune shifts. It keeps fixed the ratio between b_direct_ref and b_cross_ref. - :param tune_shifts: the sequence of complex tune shifts - :param q_s: the synchrotron tune - :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in - the x plane or $\alpha_y \sigma_y$ if working in the y plane) - :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in - the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) - :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' - :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the - stability diagram), one can consider positive tuneshift up to this fraction of q_s (default=5%) - :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding - and for the final check that the roots are actually proper roots. - :return: the detuning coefficients corresponding to the most stringent stability diagram threshold for all the - given tune shifts if the corresponding mode is unstable, 0 if all modes are stable or np.nan if the - no threshold can be found (failure of Newton's algorithm). - """ - # find max octupole current required from a list of modes, given their tuneshifts - b_coefficients = np.array([find_detuning_coeffs_threshold( - tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct_ref, - b_cross_ref=b_cross_ref, distribution=distribution, - fraction_of_qs_allowed_on_positive_side=fraction_of_qs_allowed_on_positive_side, - tolerance=tolerance) for tune_shift in tune_shifts if tune_shift is not np.nan]) - - return max(b_coefficients, key=abs_first_item_or_nan) +from xwakes.wit.landau_damping import * diff --git a/pywit/materials.py b/pywit/materials.py index 8fba2858..0441f038 100644 --- a/pywit/materials.py +++ b/pywit/materials.py @@ -1,240 +1 @@ -from pywit.interface import Layer -from pywit.utils import round_sigfigs - -from pathlib import Path -import numpy as np -import json -from typing import Callable, Tuple - - -def layer_from_dict(thickness: float, material_dict: dict) -> Layer: - """ - Define a layer from a dictionary containing the materials properties. - - :param thickness: layer thickness in m - :type thickness: float - :param material_dict: dictionary of material properties. 'dc_resistivity', 'resistivity_relaxation_time', - 're_dielectric_constant', 'magnetic_susceptibility', 'permeability_relaxation_frequency' are required - :type material_dict: dict - :return: Layer of the provided material - :rtype: Layer - """ - - # Check that the provided dict has the required entries to create a Layer object - # If not raise an AssertionError and indicate which property is missing - required_material_properties = np.array(['dc_resistivity', 'resistivity_relaxation_time', - 're_dielectric_constant', 'magnetic_susceptibility', - 'permeability_relaxation_frequency']) - # missing_properties_list is an array of bool. False indicates a missing key - missing_properties_list = np.array([key not in material_dict for key in required_material_properties]) - - assert not any(missing_properties_list), '{} missing from the input dictionary'.format( - ", ".join(required_material_properties[np.asarray(missing_properties_list)])) - - return Layer(thickness=thickness, - dc_resistivity=material_dict['dc_resistivity'], - resistivity_relaxation_time=material_dict['resistivity_relaxation_time'], - re_dielectric_constant=material_dict['re_dielectric_constant'], - magnetic_susceptibility=material_dict['magnetic_susceptibility'], - permeability_relaxation_frequency=material_dict['permeability_relaxation_frequency']) - - -def layer_from_json_material_library(thickness: float, material_key: str, - library_path: Path = Path(__file__).parent.joinpath('materials.json')) -> Layer: - """ - Define a layer using the materials.json library of materials properties. - - :param thickness: layer thickness in m - :type thickness: float - :param material_key: material key in the materials.json file - :type material_key: str - :param library_path: material library path, defaults to materials.json present in pywit - :type library_path: Path, optional - :return: Layer of the selected material - :rtype: Layer - """ - - materials_library_json = json.loads(library_path.read_bytes()) - - # Check that the material is in the provided materials library - assert material_key in materials_library_json.keys(), f"Material {material_key} is not in library {library_path}" - - # Put the material properties in a dict - # The decoder will have converted NaN, Infinity and -Infinity from JSON to python nan, inf and -inf - # The entries in the materials library can contain additional fields (comment, reference) that are not - # needed for the Layer object creation. - material_properties_dict = materials_library_json[material_key] - - layer = layer_from_dict(thickness=thickness, material_dict=material_properties_dict) - - return layer - - -# Resistivity rho at B=0 vs. temperature and RRR -def rho_vs_T_Hust_Lankford(T: float, rho273K: float, RRR: float, - P: Tuple[float, float, float, float, float, float, float], - rhoc: Callable[[float], float] = lambda T: 0) -> float: - """ - Define a resistivity versus temperature and RRR law as found in Hust & Lankford, "Thermal - conductivity of aluminum, copper, iron, and tungsten for temperatures from 1K to the melting point", National Bureau - of Standards, 1984, Eqs. (1.2.3) to (1.2.6) p. 8. - Note that typically P4 given there has the wrong sign - for Cu and W at least one would otherwise - not get the right behaviour at high temperature - most probably there is a typo in the reference. - - :param T: temperature in K - :type T: float - :param rho273K: resistivity at 273 K in Ohm.m - :type rho273K: float - :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0 - :type RRR: float - :param P: tuple of the fitting coeeficient. The coefficient can be found in the reference above for copper (p. 22), - aluminum (p. 92), iron (p. 145) and tungsten (p. 204) - :type P: tuple - :param rhoc: function of temperature accounting for the residual deviation from the law in Ohm.m, defaults to 0. - This can be used with iron and tungsten for which the residual function is provided in the reference - :type rhoc: float - :return: the resistivity value, in Ohm.m - :rtype: float - """ - - # To follow the notation used in the book (coeficients from P1 to P7), we unpack the P tuple to new variables - P1 = P[0] - P2 = P[1] - P3 = P[2] - P4 = P[3] - P5 = P[4] - P6 = P[5] - P7 = P[6] - rho0 = rho273K/RRR - rhoi = P1*T**P2 / (1. + P1*P3*T**(P2+P4)*np.exp(-(P5/T)**P6)) + rhoc(T) - rhoi0 = P7*rhoi*rho0 / (rhoi+rho0) - - return (rho0+rhoi+rhoi0) - - -# magnetoresistance law (Kohler-like) drho/rho = f(B*rho_273K(B=0)/rho_T(B=0)) -def magnetoresistance_Kohler(B_times_Sratio: float, P: Tuple[float, ...]) -> float: - """ - Define a magnetoresistance law in the form drho/rho = f(B*rho_273K(B=0)/rho_T(B=0)) - - :param B_times_Sratio: product of magnetic field B and Sratio = rho_273K(B=0)/rho_T(B=0) at a given temperature T, in Tesla - :type B_times_Sratio: float - :param P: tuple of the Kohler curve fitting coefficients. Kohler curve are represented in log-log scales: if x and y are read on the - curve, the P coefficients come from the fitting of log10(x) and log10(y) - :type P: tuple of floats - :return: drho/rho the resistivity variation at a given magnetic field - :rtype: float - """ - - if B_times_Sratio == 0. or P == (0,): - return 0. - else: - return 10.**np.polyval(P, np.log10(B_times_Sratio)) - - -def copper_at_temperature(thickness: float, T: float = 300, RRR: float = 70, B: float = 0) -> Layer: - """ - Define a layer of pure copper material at any temperature, any B field and any RRR. - We use a magnetoresistance law fitted from the UPPER curve of the plot in NIST, "Properties of copper and copper - alloys at cryogenic temperatures", by Simon, Crexler and Reed, 1992 (p. 8-27, Fig. 8-14). - The upper curve was chosen, as it is in relative agreement with C. Rathjen measurements and actual LHC beam screens - (CERN EDMS document Nr. 329882). - The resistivity vs. temperature and RRR is found from Hust & Lankford (see above). - The law vs. temperature was found in good agreement with the NIST reference above (p. 8-5, Fig. 8-2). - - :param thickness: material thickness in m - :type thickness: float - :param T: temperature in K, defaults to 300 - :type T: float, optional - :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0, defaults to 70 - :type RRR: float, optional - :param B: magnetic field in T, defaults to 0 - :type B: float, optional - :return: a Layer object - :rtype: Layer - """ - - rho273K = 15.5*1e-9 # resistivity at 273K, in Ohm.m - - # resistivity vs temperature law coefficients, found in p. 22 of Hust and Lankford - # Here P4 = -1.14 (instead of +1.14 in the book) to get the proper law behavior - P = (1.171e-17, 4.49, 3.841e10, -1.14, 50., 6.428, 0.4531) - - # Coefficients for the magnetoresistance law, from Simon, Crexler and Reed (p. 8-27) - kohler_P = (0.029497104404715, 0.905633738689341, -2.361415783729567) - - rhoDC_B0 = rho_vs_T_Hust_Lankford(T, rho273K, RRR, P) # resistivity for B=0 - Sratio = rho_vs_T_Hust_Lankford(273, rho273K, RRR, P) / rhoDC_B0 - dc_resistivity = round_sigfigs(rhoDC_B0 * (1.+magnetoresistance_Kohler(B*Sratio, kohler_P)),3) # we round it to 3 significant digits - - # tauAC formula from Ascroft-Mermin (Z=1 for copper), Drude model (used also for other - # materials defined above) with parameters from CRC - Handbook of Chem. and Phys. - me = 9.10938e-31 # electron mass - e = 1.60218e-19 # electron elementary charge - rho_m = 9.0 # Cu volumic mass in g/cm3 (it is actually 9.02 at 4K, 8.93 at 273K, - # see https://www.copper.org/resources/properties/cryogenic/ ) - A = 63.546 # Cu atomic mass in g/mol - Z = 1 # number of valence electrons - n = 6.022e23*Z*rho_m*1e6 / A - tauAC = round_sigfigs(me / (n*dc_resistivity*e**2),3) # relaxation time (s) (3 significant digits) - - return Layer(thickness=thickness, - dc_resistivity=dc_resistivity, - resistivity_relaxation_time=tauAC, - re_dielectric_constant=1, - magnetic_susceptibility=0, - permeability_relaxation_frequency=np.inf) - - -def tungsten_at_temperature(thickness: float, T: float = 300, RRR: float = 70, B: float = 0) -> Layer: - """ - Define a layer of tungsten at any temperature, any B field and any RRR. - The resistivity vs. temperature and RRR is found from Hust & Lankford (see above). - The magnetoresistance effect is not included yet. - - :param thickness: material thickness in m - :type thickness: float - :param T: temperature in K, defaults to 300 - :type T: float, optional - :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0, defaults to 70 - :type RRR: float, optional - :param B: magnetic field in T, defaults to 0 - :type B: float, optional - :return: a Layer object - :rtype: Layer - """ - - rho273K = 48.4*1e-9 # resistivity at 273K, in Ohm.m - - # resistivity vs temperature law coefficients, found in p. 204 of Hust and Lankford - # Here P4 = -1.22 (instead of +1.22 in the book) to get the proper law behavior - P = (4.801e-16, 3.839, 1.88e10, -1.22, 55.63, 2.391, 0.0) - - # Residual equation for tungsten is found in p. 204 - # The correction is small and affects only the second or third decimal of the resistivity. - # We must therefore increase the number of decimal to see its effect - rhoc = lambda T: 0.7e-8 * np.log(T/560) * np.exp(-(np.log(T/1000)/0.6)**2) - - # Coefficients for the magnetoresistance law. Put to zero for now - kohler_P = (0,) - - rhoDC_B0 = rho_vs_T_Hust_Lankford(T, rho273K, RRR, P, rhoc) # resistivity for B=0 - Sratio = rho_vs_T_Hust_Lankford(273, rho273K, RRR, P, rhoc) / rhoDC_B0 - dc_resistivity = round_sigfigs(rhoDC_B0 * (1.+magnetoresistance_Kohler(B*Sratio, kohler_P)),3) # 3 significant digits - - # tauAC formula from Ascroft-Mermin (Z=2 for tungsten), Drude model (used also for other - # materials defined above) with parameters from CRC - Handbook of Chem. and Phys. - me = 9.10938e-31 # electron mass - e = 1.60218e-19 # electron elementary charge - rhom = 19.3 # W volumic mass in g/cm3 - A = 183.84 # W atomic mass in g/mol - Z = 2 # number of valence electrons - n = 6.022e23 * Z * rhom * 1e6 / A - tauAC = round_sigfigs(me/(n*dc_resistivity*e**2),3) # relaxation time (s) (3 significant digits) - - return Layer(thickness=thickness, - dc_resistivity=dc_resistivity, - resistivity_relaxation_time=tauAC, - re_dielectric_constant=1, - magnetic_susceptibility=0, - permeability_relaxation_frequency=np.inf) +from xwakes.wit.materials import * diff --git a/pywit/model.py b/pywit/model.py index 1f2c9a8c..08d74639 100644 --- a/pywit/model.py +++ b/pywit/model.py @@ -1,29 +1 @@ -from pywit.element import Element - -from typing import List, Optional, Tuple - - -class Model: - """ - Suggestion for structure of Model class - """ - def __init__(self, elements: List[Element] = None, lumped_betas: Optional[Tuple[float, float]] = None): - assert elements, "Model object needs to be initialized with at least one Element" - if lumped_betas is not None: - elements = [element.changed_betas(*lumped_betas) for element in elements] - - self.__elements = elements - self.__lumped_betas = lumped_betas - - @property - def elements(self): - return self.__elements - - @property - def total(self): - return sum(self.__elements) - - def append_element(self, element: Element): - if self.__lumped_betas is not None: - element = element.changed_betas(*self.__lumped_betas) - self.__elements.append(element) +from xwakes.wit.model import * diff --git a/pywit/parameters.py b/pywit/parameters.py index 8dfe3b8f..2e3c6dc4 100644 --- a/pywit/parameters.py +++ b/pywit/parameters.py @@ -1,11 +1 @@ -REL_TOL = 1e-8 -ABS_TOL = 1e-8 -PRECISION_FACTOR = 100. - -MIN_FREQ = 1e3 -MAX_FREQ = 1e10 -FREQ_P_FACTOR = 0.2 - -MIN_TIME = 1e-10 -MAX_TIME = 1e-7 -TIME_P_FACTOR = 0.2 +from xwakes.wit.parameters import * diff --git a/pywit/plot.py b/pywit/plot.py index dce551c5..bc606669 100644 --- a/pywit/plot.py +++ b/pywit/plot.py @@ -1,230 +1 @@ -from pywit.component import Component -from pywit.element import Element -from pywit.model import Model -from pywit.parameters import * - -from typing import List, Dict, Union, Optional, Set -from collections import defaultdict - -import matplotlib.pyplot as plt -import numpy as np - - -def plot_component(component: Component, plot_impedance: bool = True, plot_wake: bool = True, start: float = 1, - stop: float = 10000, points: int = 200, step_size: float = None, plot_real: bool = True, - plot_imag: bool = True) -> None: - """ - Function for plotting real and imaginary parts of impedance and wake functions of a single component - :param component: The component to be plotted - :param plot_impedance: A flag indicating if impedance should be plotted - :param plot_wake: A flag indicating if wake should be plotted - :param start: The first value on the x-axis of the plot - :param stop: The last value on the x-axis of the plot - :param points: The number of points to be evaluated for each line of the plot (alternative to step_size) - :param step_size: The distance between each point on the x-axis (alternative to points) - :param plot_real: A flag indicating if the real values should be plotted - :param plot_imag: A flag indicating if the imaginary values should be plotted - :return: Nothing - """ - assert (plot_wake or plot_impedance) and (plot_real or plot_imag), "There is nothing to plot" - assert stop - start > 0, "stop must be greater than start" - if step_size: - assert step_size > 0, "Negative step_size not possible" - xs = np.arange(start, stop, step_size, dtype=float) - else: - xs = np.linspace(start, stop, points) - - legends = [] - if plot_impedance: - ys = component.impedance(xs) - if plot_real: - plt.plot(xs, ys.real) - legends.append("Re[Z(f)]") - if plot_imag: - plt.plot(xs, ys.imag) - legends.append("Im[Z(f)]") - if plot_wake: - ys = component.wake(xs) - if plot_real: - plt.plot(xs, ys.real) - legends.append("Re[W(z)]") - if plot_imag: - plt.plot(xs, ys.imag) - legends.append("Im[W(z)]") - - plt.legend(legends) - plt.show() - - -def plot_element_in_plane(element: Element, plane: str, plot_impedance: bool = True, plot_wake: bool = True, - start: float = 1, stop: float = 10000, points: int = 200, step_size: float = None, - plot_real: bool = True, plot_imag: bool = True): - """ - Sums all Components of Element in specified plane and plots the resulting Component with the parameters given. - :param element: The Element with Components to be plotted - :param plane: The plane in which the Components to be plotted lie - :param plot_impedance: A flag indicating if impedance should be plotted - :param plot_wake: A flag indicating if wake should be plotted - :param start: The first value on the x-axis of the plot - :param stop: The last value on the x-axis of the plot - :param points: The number of points to be evaluated for each line of the plot (alternative to step_size) - :param step_size: The distance between each point on the x-axis (alternative to points) - :param plot_real: A flag indicating if the real values should be plotted - :param plot_imag: A flag indicating if the imaginary values should be plotted - :return: Nothing - """ - plane = plane.lower() - assert plane in ['x', 'y', 'z'], f"{plane} is not a valid plane. Must be 'x', 'y' or 'z'" - component = sum([c for c in element.components if c.plane == plane]) - - assert component, f"Element has no components in plane {plane}" - plot_component(component, plot_impedance=plot_impedance, plot_wake=plot_wake, start=start, stop=stop, - points=points, step_size=step_size, plot_real=plot_real, plot_imag=plot_imag) - - -def plot_component_impedance(component: Component, logscale_x: bool = True, logscale_y: bool = True, - points: int = 1000, start=MIN_FREQ, stop=MAX_FREQ, title: Optional[str] = None) -> None: - fig: plt.Figure = plt.figure() - ax: plt.Axes = fig.add_subplot(111) - fs = np.geomspace(start, stop, points) - impedances = component.impedance(fs) - reals, imags = impedances.real, impedances.imag - ax.plot(fs, reals, label='real') - ax.plot(fs, imags, label='imag') - ax.set_xlim(start, stop) - if title: - plt.title(title) - - if logscale_x: - ax.set_xscale('log') - - if logscale_y: - ax.set_yscale('log') - - ax.set_xlabel('Frequency [Hz]') - ax.set_ylabel('Z [Ohm / m]') - - plt.legend() - plt.show() - - -def plot_component_wake(component: Component, logscale_x: bool = True, logscale_y: bool = True, - points: int = 1000, start=MIN_TIME, stop=MAX_TIME, title: Optional[str] = None) -> None: - fig: plt.Figure = plt.figure() - ax: plt.Axes = fig.add_subplot(111) - ts = np.geomspace(start, stop, points) - ax.plot(ts, component.wake(ts)) - ax.set_xlim(start, stop) - - ax.set_xlabel('Time [s]') - ax.set_ylabel('Wake [V / C * m]') - - if title: - plt.title(title) - - if logscale_x: - ax.set_xscale('log') - - if logscale_y: - ax.set_yscale('log') - - plt.show() - - -def generate_contribution_plots(model: Model, start_freq: float = MIN_FREQ, stop_freq: float = MAX_FREQ, - start_time: float = MIN_TIME, stop_time: float = MAX_TIME, points: int = 1000, - freq_scale: str = 'log', time_scale: str = 'log', absolute: bool = False) -> None: - # TODO: use roi's to generate grid - fs = np.geomspace(start_freq, stop_freq, points) - ts = np.geomspace(start_time, stop_time, points) - - all_tags = set([e.tag for e in model.elements]) - elements: Dict[str, Union[int, Element]] = defaultdict(int) - - for element in model.elements: - elements[element.tag] += element - - components_defined_for_tag: Dict[str, Set[str]] = dict() - all_type_strings = set() - for tag, element in elements.items(): - components_defined_for_tag[tag] = {c.get_shorthand_type() for c in element.components} - all_type_strings.update(components_defined_for_tag[tag]) - - tags = list(all_tags) - cumulative_elements: List[Element] = [] - defined_type_strings: List[set] = [] - current_defined = set() - current_element = 0 - for tag in tags: - new_element = current_element + elements[tag] - current_defined = current_defined.union(components_defined_for_tag[tag]) - defined_type_strings.append(current_defined) - cumulative_elements.append(new_element) - current_element = new_element - - for type_string in all_type_strings: - wakes = np.asarray([np.zeros(shape=ts.shape)]) - real_impedances = np.asarray([np.zeros(shape=fs.shape)]) - imag_impedances = np.asarray([np.zeros(shape=fs.shape)]) - for i, element in enumerate(cumulative_elements): - if type_string not in defined_type_strings[i]: - wakes = np.vstack((wakes, wakes[-1])) - real_impedances = np.vstack((real_impedances, real_impedances[-1])) - imag_impedances = np.vstack((imag_impedances, imag_impedances[-1])) - continue - - component = [c for c in element.components if c.get_shorthand_type() == type_string][0] - if component.wake is not None: - array = component.wake(ts) - wakes = np.vstack((wakes, array)) - else: - wakes = np.vstack((wakes, wakes[-1])) - - if component.impedance is not None: - impedances = component.impedance(fs) - real_impedances = np.vstack((real_impedances, impedances.real)) - imag_impedances = np.vstack((imag_impedances, impedances.imag)) - else: - real_impedances = np.vstack((real_impedances, real_impedances[-1])) - imag_impedances = np.vstack((imag_impedances, imag_impedances[-1])) - - titles = (f'Re[Z] contribution - {type_string}', - f'Im[Z] contribution - {type_string}', - f'Wake contribution - {type_string}') - - for i, (array, xs, title) in enumerate(zip((real_impedances, imag_impedances, wakes), - (fs, fs, ts), titles)): - if sum(array[-1]) == 0: - continue - - if not absolute: - array = np.divide(array, array[-1]) - - fig: plt.Figure = plt.figure() - ax: plt.Axes = fig.add_subplot(111) - for j in range(len(all_tags) - 1, -1, -1): - ax.fill_between(xs, array[j], array[j + 1], label=tags[j]) - - ax.set_xscale(time_scale if i == 2 else freq_scale) - if not absolute: - ax.set_ylim(0, 1) - ax.set_xlim(start_time if i == 2 else start_freq, stop_time if i == 2 else stop_freq) - plt.title(title) - plt.legend() - plt.show() - - -def plot_total_impedance_and_wake(model: Model, logscale_y=True, start_freq: float = MIN_FREQ, - stop_freq: float = MAX_FREQ, start_time: float = MIN_TIME, - stop_time: float = MAX_TIME, points: int = 1000, - logscale_freq: bool = True, logscale_time: bool = True): - element = sum(model.elements) - for component in element.components: - if component.impedance: - plot_component_impedance(component, logscale_y=logscale_y, - title=f"Total impedance - {component.get_shorthand_type()}", start=start_freq, - stop=stop_freq, points=points, logscale_x=logscale_freq) - if component.wake: - plot_component_wake(component, logscale_y=logscale_y, - title=f"Total wake - {component.get_shorthand_type()}", start=start_time, - stop=stop_time, points=points, logscale_x=logscale_time) +from xwakes.wit.plot import * diff --git a/pywit/sacherer_formula.py b/pywit/sacherer_formula.py index 194814fe..de7b4779 100644 --- a/pywit/sacherer_formula.py +++ b/pywit/sacherer_formula.py @@ -1,199 +1 @@ -import numpy as np -import sys - -from scipy.constants import e as q_p, m_p, c - -from typing import List, Callable, Iterable, Tuple, Union - - -def sacherer_formula(qp: float, nx_array: np.array, bunch_intensity: float, omegas: float, n_bunches: int, - omega_rev: float, tune: float, gamma: float, eta: float, bunch_length_seconds: float, m_max: int, - impedance_table: np.array = None, freq_impedance_table: np.array = None, - impedance_function: Callable[[float], float] = None, m0: float = m_p, charge: float = q_p, - mode_type: str = 'sinusoidal') -> Tuple[np.array, float, np.array]: - """ - Computes frequency shift and effective impedance from Sacherer formula, in transverse, in the case of low - intensity perturbations (no mode coupling), for modes of kind 'mode_type'. - - Documentation: see Elias Metral's USPAS 2009 course : Bunched beams transverse coherent - instabilities. - - NOTE: this is NOT the original Sacherer formula, which assumes an impedance normalized by beta - (see E. Metral, USPAS 2009 lectures, or C. Zannini, - https://indico.cern.ch/event/766028/contributions/3179810/attachments/1737652/2811046/Z_definition.pptx) - Here this formula is instead divided by beta (compared to Sacherer initial one), - so is valid with our usual definition of impedance (not beta-normalized). - - :param qp: the chromaticity (defined as $\frac{\Delta q \cdot p}{\Delta p}$ - :param nx_array: a vector of coupled bunch modes for which the tune shift is computed (it must contain integers in - the range (0, M-1)) - :param bunch_intensity: number of particles per bunch - :param omegas: the synchrotron angular frequency (i.e. $Q_s \cdot \omega_{rev}$) - :param n_bunches: the number of bunches - :param omega_rev: the revolution angular frequency (i.e. $2\cdot \pi f_{rev}$) - :param tune: machine tune in the considered plane (the TOTAL tune, including the integer part, must be passed) - :param gamma: the relativistic gamma factor of the beam - :param eta: the slippage factor (i.e. alpha_p - 1/gamma^2, with alpha_p the momentum compaction factor) - :param bunch_length_seconds: the total bunch length in seconds (4 times $\sigma$ for a Gaussian bunch) - :param m_max: specifies the range (-m_max to m_max) of the azimuthal modes to be considered - :param impedance_table: a numpy array giving the complex impedance at a discrete set of points. It must be specified - if impedance_function is not specified - :param freq_impedance_table: the frequencies at which the impedance is sampled. It must be specified if - impedance_function is not specified - :param impedance_function: the impedance function. It must be specified if impedance_table is not specified - :param m0: the rest mass of the considered particles - :param charge: the charge of the considered particles - :param mode_type: the type of modes in the effective impedance. It can be 'sinusoidal' (typically - well-adpated for protons) or 'hermite' (typically better for leptons). - - :return tune_shift_nx: tune shifts for all multibunch modes and synchrotron modes. It is an array of dimensions - ( len(nx_scan), (2*m_max+1) ) - :return tune_shift_m0: tune shift of the most unstable coupled-bunch mode with m=0 - :return effective_impedance: the effective impedance for all multibunch modes and synchrotron modes. It is an array - of dimensions ( len(nx_scan), (2*m_max+1) ) - """ - - def hmm(m_mode: int, omega: Union[float, np.ndarray]): - """ - Compute hmm power spectrum of Sacherer formula, for azimuthal mode number m, - at angular frequency 'omega' (rad/s) (can be an array), for total bunch length - 'bunch_length_seconds' (s), and for a kind of mode specified by 'mode_type' - (which can be 'hermite' or 'sinusoidal') - :param m_mode: the azimuthal mode number - :param omega: the angular frequency at which hmm is computed - """ - - if mode_type.lower().startswith('sinus'): - # best for protons - hmm_val = (((bunch_length_seconds * (np.abs(m_mode) + 1.)) ** 2 / (2. * np.pi ** 4)) * - (1. + (-1) ** m_mode * np.cos(omega * bunch_length_seconds)) / - (((omega * bunch_length_seconds / np.pi) ** 2 - (np.abs(m_mode) + 1.) ** 2) ** 2)) - - elif mode_type.lower() == 'hermite': - # best for leptons - hmm_val = (omega * bunch_length_seconds / 4) ** (2 * m_mode) * np.exp( - -(omega * bunch_length_seconds / 4.) ** 2) - - else: - raise ValueError("mode_type can only be 'sinusoidal' or 'hermite'") - - return hmm_val - - def hmm_weighted_sum(m_mode: int, nx_mode: int, weight_function: Callable[[float], complex] = None): - """ - Compute sum of hmm functions in the Sacherer formula, optionally - weighted by weight_function. - Note: In the end the sum runs over k with hmm taken at the angular frequencies - (k_offset+k*n_bunches)*omega0+m*omegas-omegaksi but the impedance is taken at - (k_offset+k*n_bunches)*omega0+m*omegas - :param m_mode: the azimuthal mode number - :param nx_mode: the coupled-bunch mode number - :param weight_function: function of frequency (NOT angular) giving - the sum weights (typically, it is the impedance) (optional) - :return: the (possibly weigthed) sum of hmm functions - """ - eps = 1.e-5 # relative precision of the summations - k_max = 20 - k_offset = nx_mode + fractional_tune - # sum initialization - omega_k = k_offset * omega_rev + m_mode * omegas - hmm_k = hmm(m_mode, omega_k - omega_ksi) - - omega = np.arange(-100.01 / bunch_length_seconds, 100.01 / bunch_length_seconds, - 0.01 / bunch_length_seconds) - - if weight_function is not None: - z_pk = weight_function(omega_k / (2 * np.pi)) - else: - z_pk = np.ones_like(omega_k) - - sum1_inner = z_pk * hmm_k - - k = np.arange(1, k_max + 1) - old_sum1 = 10. * sum1_inner - - while ((np.abs(np.real(sum1_inner - old_sum1))) > eps * np.abs(np.real(sum1_inner))) or ( - (np.abs(np.imag(sum1_inner - old_sum1))) > eps * np.abs(np.imag(sum1_inner))): - old_sum1 = sum1_inner - # omega_k^x and omega_-k^x in Elias's slides: - omega_k = (k_offset + k * n_bunches) * omega_rev + m_mode * omegas - omega_mk = (k_offset - k * n_bunches) * omega_rev + m_mode * omegas - # power spectrum function h(m,m) for k and -k: - hmm_k = hmm(m_mode, omega_k - omega_ksi) - hmm_mk = hmm(m_mode, omega_mk - omega_ksi) - - if weight_function is not None: - z_pk = weight_function(omega_k / (2 * np.pi)) - z_pmk = weight_function(omega_mk / (2 * np.pi)) - else: - z_pk = np.ones_like(omega_k) - z_pmk = np.ones_like(omega_mk) - - # sum - sum1_inner = sum1_inner + np.sum(z_pk * hmm_k) + np.sum(z_pmk * hmm_mk) - - k = k + k_max - - return sum1_inner - - if impedance_function is not None and impedance_table is not None: - raise ValueError('Only one between impedance_function and impedance_table can be specified') - - if impedance_table is not None and freq_impedance_table is None: - raise ValueError('When impedance_table is specified, also the corresponding frequencies must be specified in' - 'omega_impedance_table') - - # some parameters - beta = np.sqrt(1. - 1. / (gamma ** 2)) # relativistic velocity factor - f0 = omega_rev / (2. * np.pi) # revolution angular frequency - single_bunch_current = charge * bunch_intensity * f0 # single-bunch current - fractional_tune = tune - np.floor(tune) # fractional part of the tune - bunch_length_seconds_meters = bunch_length_seconds * beta * c # full bunch length (in meters) - - if impedance_table is not None: - def impedance_function(x): - if np.isscalar(x): - x = np.array([x]) - ind_p = x >= 0 - ind_n = x < 0 - result = np.zeros_like(x, dtype=complex) - result[ind_p] = np.interp(x[ind_p], freq_impedance_table, impedance_table) - result[ind_n] = -np.interp(np.abs(x[ind_n]), freq_impedance_table, impedance_table).conjugate() - - return result - - tune_shift_nx = np.zeros((len(nx_array), 2 * m_max + 1), dtype=complex) - tune_shift_m0 = complex(0) - effective_impedance = np.zeros((len(nx_array), 2 * m_max + 1), dtype=complex) - - omega_ksi = qp * omega_rev / eta - - for inx, nx in enumerate(nx_array): # coupled-bunch modes - - for im, m in enumerate(range(-m_max, m_max + 1)): - # consider each synchrotron mode individually - # sum power spectrum functions and computes effective impedance - - # sum power functions - # BE CAREFUL: maybe for this "normalization sum" the sum should run - # on all single-bunch harmonics instead of only coupled-bunch - # harmonics (and then the frequency shift should be multiplied by - # n_bunches). This has to be checked. - sum1 = hmm_weighted_sum(m, nx) - - # effective impedance - sum2 = hmm_weighted_sum(m, nx, weight_function=impedance_function) - - effective_impedance[inx, im] = sum2 / sum1 - freq_shift = 1j * charge * single_bunch_current / (2 * (np.abs(m) + 1.) * m0 * gamma * tune * omega_rev * - bunch_length_seconds_meters) * sum2 / sum1 - - tune_shift = (freq_shift / omega_rev + m * omegas / omega_rev) - - tune_shift_nx[inx, im] = tune_shift - - # find the most unstable coupled-bunch mode for m=0 - inx = np.argmin(np.imag(tune_shift_nx[:, m_max])) - tune_shift_m0 = tune_shift_nx[inx, m_max] - - return tune_shift_nx, tune_shift_m0, effective_impedance +from xwakes.wit.sacherer_formula import * diff --git a/pywit/utilities.py b/pywit/utilities.py index 8c32d410..f61c1064 100644 --- a/pywit/utilities.py +++ b/pywit/utilities.py @@ -1,727 +1 @@ -from pywit.component import Component -from pywit.element import Element -from pywit.interface import Layer, FlatIW2DInput, RoundIW2DInput -from pywit.interface import component_names - -from yaml import load, SafeLoader -from typing import Tuple, Dict, List, Union, Sequence, Optional, Callable -from collections import defaultdict - -from numpy import (vectorize, sqrt, exp, pi, sin, cos, abs, sign, - inf, floor, linspace, trapz, ones, isscalar, array) -from numpy.typing import ArrayLike -import scipy.constants -from scipy import special as sp -import numpy as np - -c_light = scipy.constants.speed_of_light # m s-1 -mu0 = scipy.constants.mu_0 -Z0 = mu0 * c_light -eps0 = scipy.constants.epsilon_0 - - -def string_to_params(name: str, include_is_impedance: bool = True): - """ - Converts a string describing some specific component to the set of parameters which give this component - :param name: The string description of the component to be converted. Either of the form where p is either - 'x', 'y' or 'z', and abcd are four integers giving the source- and test exponents of the component. Optionally, - a single character, 'z' or 'w' can be included at the front of the string if it is necessary to indicate that the - given component is either a wake- or impedance component. In this case, the include_is_impedance flag must be set. - :param include_is_impedance: A flag indicating whether or not "name" includes a 'z' or 'w' at the beginning. - :return: A tuple containing the necessary information to describe the type of the given component - """ - is_impedance = False - if include_is_impedance: - is_impedance = name[0] == 'z' - name = name[1:] - - plane = name[0].lower() - exponents = tuple(int(n) for n in name[1:]) - - return (is_impedance, plane, exponents) if include_is_impedance else (plane, exponents) - - -def create_component_from_config(identifier: str) -> Component: - """ - Creates a Component object as specified by a .yaml config file - :param identifier: A unique identifier-string corresponding to a component specification in config/component.yaml - :return: A Component object initialized according to the specification - """ - with open("config/component.yaml", "r") as f: - cdict = load(f, Loader=SafeLoader)[identifier] - wake = vectorize(lambda x: eval(cdict['wake'], {'exp': exp, 'sqrt': sqrt, - 'pi': pi, 'x': x})) if 'wake' in cdict else None - impedance = vectorize(lambda x: eval(cdict['impedance'], {'exp': exp, 'sqrt': sqrt, - 'pi': pi, 'x': x})) if 'impedance' in cdict else None - name = cdict['name'] if 'name' in cdict else "" - plane = cdict['plane'] - source_exponents = (int(cdict['source_exponents'][0]), int(cdict['source_exponents'][-1])) - test_exponents = (int(cdict['test_exponents'][0]), int(cdict['test_exponents'][-1])) - return Component(impedance, wake, plane, source_exponents, test_exponents, name) - - -def create_element_from_config(identifier: str) -> Element: - """ - Creates an Element object as specified by a .yaml config file - :param identifier: A unique identifier-string corresponding to an element specification in config/element.yaml - :return: An Element object initialized according to the specifications - """ - with open("config/element.yaml", "r") as f: - edict = load(f, Loader=SafeLoader)[identifier] - length = float(edict['length']) if 'length' in edict else 0 - beta_x = float(edict['beta_x']) if 'beta_x' in edict else 0 - beta_y = float(edict['beta_y']) if 'beta_y' in edict else 0 - name = edict['name'] if 'name' in edict else "" - tag = edict['tag'] if 'tag' in edict else "" - components = [create_component_from_config(c_id) for c_id in edict['components'].split()] \ - if 'components' in edict else [] - return Element(length, beta_x, beta_y, components, name, tag) - - -def compute_resonator_f_roi_half_width(q: float, f_r: float, f_roi_level: float = 0.5): - aux = np.sqrt((1 - f_roi_level) / f_roi_level) - - return (aux + np.sqrt(aux**2 + 4*q**2))*f_r/(2*q) - f_r - - -def create_resonator_component(plane: str, exponents: Tuple[int, int, int, int], - r: float, q: float, f_r: float, f_roi_level: float = 0.5) -> Component: - """ - Creates a single component object belonging to a resonator - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param r: the shunt impedance of the given component of the resonator - :param q: the quality factor of the given component of the resonator - :param f_r: the resonance frequency of the given component of the resonator - :param f_roi_level: fraction of the peak ok the resonator which is covered by the ROI. I.e. the roi will cover - the frequencies for which the resonator impedance is larger than f_roi_level*r - :return: A component object of a resonator, specified by the input arguments - """ - root_term = sqrt(1 - 1 / (4 * q ** 2) + 0J) - omega_r = 2 * pi * f_r - if plane == 'z': - impedance = lambda f: r / (1 - 1j * q * (f_r / f - f / f_r)) - omega_bar = omega_r * root_term - alpha = omega_r / (2 * q) - wake = lambda t: (omega_r * r * exp(-alpha * t) * ( - cos(omega_bar * t) - alpha * sin(omega_bar * t) / omega_bar) / q).real - else: - impedance = lambda f: (f_r * r) / (f * (1 - 1j * q * (f_r / f - f / f_r))) - wake = lambda t: (omega_r * r * exp(-omega_r * t / (2 * q)) * sin(omega_r * root_term * t) / - (q * root_term)).real - - if q > 1: - d = compute_resonator_f_roi_half_width(q=q, f_r=f_r, f_roi_level=f_roi_level) - f_rois = [(f_r - d, f_r + d)] - else: - f_rois = [] - # TODO: add ROI(s) for wake - - return Component(impedance, wake, plane, source_exponents=exponents[:2], - test_exponents=exponents[2:], - f_rois=f_rois) - - -def create_resonator_element(length: float, beta_x: float, beta_y: float, - rs: Dict[str, float], qs: Dict[str, float], fs: Dict[str, float], - f_roi_levels: Dict[str, float] = None, tag: str = 'resonator', - description: str = '') -> Element: - """ - Creates an element object representing a resonator. - :param length: The length, in meters, of the resonator element - :param beta_x: The value of the beta function in the x-direction at the position of the resonator element - :param beta_y: The value of the beta function in the y-direction at the position of the resonator element - :param rs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the - values give the Shunt impedance corresponding to this particular component - :param qs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the - values give the quality factor of the specified component of the resonator - :param fs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the - values give the resonance frequency corresponding to the particular component - :param f_roi_levels: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and - the values give the fraction of the peak ok the resonator which is covered by the ROI. I.e. the roi will cover - the frequencies for which the resonator impedance is larger than f_roi_level*r - :param tag: An optional short string used to place elements into categories - :param description: An optional short description of the element - :return: An element object as specified by the user-input - """ - if f_roi_levels is None: - f_roi_levels = {} - for key in rs.keys(): - f_roi_levels[key] = 0.5 - - assert set(rs.keys()) == set(qs.keys()) == set(fs.keys()) == set(f_roi_levels.keys()), "The three input " \ - "dictionaries describing " \ - "the resonator do not all " \ - "have identical keys" - components = [] - for key in rs.keys(): - plane, exponents = string_to_params(key, include_is_impedance=False) - components.append(create_resonator_component(plane, exponents, rs[key], qs[key], fs[key], - f_roi_level=f_roi_levels[key])) - - return Element(length, beta_x, beta_y, components, tag=tag, description=description) - - -def create_many_resonators_element(length: float, beta_x: float, beta_y: float, - params_dict: Dict[str, List[Dict[str, float]]], tag: str = 'resonator', - description: str = '') -> Element: - """ - Creates an element object representing many resonators. - :param length: The length, in meters, of the element - :param beta_x: The value of the beta function in the x-direction at the position of the element - :param beta_y: The value of the beta function in the y-direction at the position of the element - :param params_dict: a dictionary associating to each component a list of dictionaries containing the - parameters of a resonator component. E.g.: - params_dict = { - 'z0000': - [ - {'r': 10, 'q': 10, 'f': 50, 'f_roi_level: 0.5}, - {'r': 40, 'q': 100, 'f': 60} - ], - 'x1000': - ... - } - f_roi_level is optional - :param tag: An optional short string used to place elements into categories - :param description: An optional short description of the element - :return: An element object as specified by the user-input - """ - all_components = [] - for component_id, component_params_list in params_dict.items(): - plane, exponents = string_to_params(component_id, include_is_impedance=False) - for component_params in component_params_list: - assert ('r' in component_params.keys() and 'q' in component_params.keys() and - 'f' in component_params.keys()), "each of the the component dictionaries must contain r, q and f" - - f_roi_level = component_params.get('f_roi_level', 0.5) - all_components.append(create_resonator_component(plane, exponents, component_params['r'], - component_params['q'], component_params['f'], - f_roi_level=f_roi_level)) - - comp_dict = defaultdict(lambda: 0) - for c in all_components: - comp_dict[(c.plane, c.source_exponents, c.test_exponents)] += c - - components = sorted(comp_dict.values(), key=lambda x: (x.plane, x.source_exponents, x.test_exponents)) - - return Element(length, beta_x, beta_y, components, tag=tag, description=description) - - -def create_classic_thick_wall_component(plane: str, exponents: Tuple[int, int, int, int], - layer: Layer, radius: float) -> Component: - """ - Creates a single component object modeling a resistive wall impedance/wake, - based on the "classic thick wall formula" (see e.g. A. W. Chao, chap. 2 in - "Physics of Collective Beams Instabilities in High Energy Accelerators", - John Wiley and Sons, 1993). - Only longitudinal and transverse dipolar impedances are supported here. - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param layer: the chamber material, as a pywit Layer object - :param radius: the chamber radius in m - :return: A component object - """ - - # Material properties required for the skin depth computation are derived from the input Layer attributes - material_resistivity = layer.dc_resistivity - material_relative_permeability = 1. + layer.magnetic_susceptibility - material_permeability = material_relative_permeability * mu0 - - # Create the skin depth as a function offrequency and layer properties - delta_skin = lambda f: (material_resistivity / (2*pi*abs(f) * material_permeability)) ** (1/2) - - # Longitudinal impedance and wake - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: (1/2) * (1+sign(f)*1j) * material_resistivity / (pi * radius) * (1 / delta_skin(f)) - wake = lambda t: - c_light / (2*pi*radius) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(1/2)) - # Transverse dipolar impedance - elif (plane == 'x' and exponents == (1, 0, 0, 0)) or (plane == 'y' and exponents == (0, 1, 0, 0)): - impedance = lambda f: ((c_light/(2*pi*f)) * (1+sign(f)*1j) * - material_resistivity / (pi * radius**3) * - (1 / delta_skin(f))) - wake = lambda t: -c_light / (2*pi*radius**3) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(3/2)) - else: - print("Warning: resistive wall impedance not implemented for component {}{}. Set to zero".format(plane, - exponents)) - impedance = lambda f: 0 - wake = lambda f: 0 - - return Component(vectorize(impedance), vectorize(wake), plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) - - -def _zlong_round_single_layer_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius: float, length: float) -> ArrayLike: - """ - Function to compute the longitudinal resistive-wall impedance from - the single-layer, approximated formula for a cylindrical structure, - by E. Metral (see e.g. Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius: the radius of the structure (in m) - :param length: the total length of the resistive object (in m) - :return: Zlong, the longitudinal impedance at these frequencies - """ - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_long = 1j*omega*mu0*length/(2.*pi*beta**2*gamma**2) - - x1 = k*radius/gamma - x1sq = x1**2 - x2 = nu*radius - - zlong = coef_long * (sp.k0(x1)/sp.i0(x1) - 1./(x1sq*(1./2.+eps1*sp.kve(1, x2)/(x2*sp.kve(0, x2))))) - - return zlong - - -def _zdip_round_single_layer_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius: float, length: float) -> ArrayLike: - """ - Function to compute the transverse dipolar resistive-wall impedance from - the single-layer, approximated formula for a cylindrical structure, - Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius: the radius of the structure (in m) - :param length: the total length of the resistive object (in m) - :return: Zdip, the transverse dipolar impedance at these frequencies - """ - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_dip = 1j*k**2*Z0*length/(4.*pi*beta*gamma**4) - - x1 = k*radius/gamma - x1sq = x1**2 - x2 = nu*radius - - zdip = coef_dip * (sp.k1(x1)/sp.i1(x1) + 4.*beta**2*gamma**2/(x1sq*(2.+x2*sp.kve(0, x2)/(mu1*sp.kve(1, x2))))) - - return zdip - - -def create_resistive_wall_single_layer_approx_component(plane: str, exponents: Tuple[int, int, int, int], - input_data: Union[FlatIW2DInput, RoundIW2DInput]) -> Component: - """ - Creates a single component object modeling a resistive wall impedance, - based on the single-layer approximated formulas by E. Metral (see e.g. - Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param input_data: an IW2D input object (flat or round). If the input - is of type FlatIW2DInput and symmetric, we apply to the round formula the - Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, - KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, - https://cds.cern.ch/record/248630/files/p221.pdf), - while for a single plate we use those from A. Burov and V. Danilov, - PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other - kinds of asymmetric structure will raise an error. - If the input is of type RoundIW2DInput, the structure is in principle round - but the Yokoya factors put in the input will be used. - :return: A component object - """ - gamma = input_data.relativistic_gamma - length = input_data.length - - if isinstance(input_data, FlatIW2DInput): - if len(input_data.top_layers) > 1: - raise NotImplementedError("Input data can have only one layer") - yok_long = 1. - layer = input_data.top_layers[0] - radius = input_data.top_half_gap - if input_data.top_bottom_symmetry: - yok_dipx = pi**2/24. - yok_dipy = pi**2/12. - yok_quax = -pi**2/24. - yok_quay = pi**2/24. - elif input_data.bottom_half_gap == inf: - yok_dipx = 0.25 - yok_dipy = 0.25 - yok_quax = -0.25 - yok_quay = 0.25 - else: - raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " - "hence the bottom half gap must be infinite") - elif isinstance(input_data, RoundIW2DInput): - radius = input_data.inner_layer_radius - if len(input_data.layers) > 1: - raise NotImplementedError("Input data can have only one layer") - layer = input_data.layers[0] - yok_long = input_data.yokoya_factors[0] - yok_dipx = input_data.yokoya_factors[1] - yok_dipy = input_data.yokoya_factors[2] - yok_quax = input_data.yokoya_factors[3] - yok_quay = input_data.yokoya_factors[4] - else: - raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") - - # Longitudinal impedance - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: yok_long*_zlong_round_single_layer_approx( - f, gamma, layer, radius, length) - # Transverse impedances - elif plane == 'x' and exponents == (1, 0, 0, 0): - impedance = lambda f: yok_dipx*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'y' and exponents == (0, 1, 0, 0): - impedance = lambda f: yok_dipy*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'x' and exponents == (0, 0, 1, 0): - impedance = lambda f: yok_quax*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - elif plane == 'y' and exponents == (0, 0, 0, 1): - impedance = lambda f: yok_quay*_zdip_round_single_layer_approx( - f, gamma, layer, radius, length) - else: - impedance = lambda f: 0 - - return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) - - -def create_resistive_wall_single_layer_approx_element( - input_data: Union[FlatIW2DInput, RoundIW2DInput], - beta_x: float, beta_y: float, - component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), - name: str = "", tag: str = "", description: str = "") -> Element: - """ - Creates an element object modeling a resistive wall impedance, - based on the single-layer approximated formulas by E. Metral (see e.g. - Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, - https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and - Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, - https://doi.org/10.1103/PhysRevSTAB.12.084401) - :param input_data: an IW2D input object (flat or round). If the input - is of type FlatIW2DInput and symmetric, we apply to the round formula the - Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, - KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, - https://cds.cern.ch/record/248630/files/p221.pdf), - while for a single plate we use those from A. Burov and V. Danilov, - PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other - kinds of asymmetric structure will raise an error. - If the input is of type RoundIW2DInput, the structure is in principle round - but the Yokoya factors put in the input will be used. - :param beta_x: The beta function in the x-plane at the position of the element - :param beta_y: The beta function in the y-plane at the position of the element - :param component_ids: a list of components to be computed - :param name: A user-specified name for the Element - :param tag: A string to tag the Element - :param description: A description for the Element - :return: An Element object representing the structure - """ - components = [] - length = input_data.length - for component_id in component_ids: - _, plane, exponents = component_names[component_id] - components.append(create_resistive_wall_single_layer_approx_component( - plane, exponents, input_data)) - - return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, - description=description) - - -def _zlong_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius_small: float, - radius_large: float, length: float, - step_size: float = 1e-3) -> ArrayLike: - """ - Function to compute the longitudinal resistive-wall impedance for a - round taper, integrating the radius-dependent approximated formula - for a cylindrical structure (see_zlong_round_single_layer_approx above), - over the length of the taper. - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param length: the total length of the taper (in m) - :param step_size: the step size (in the radial direction) for the - integration (in m) - :return: the longitudinal impedance at these frequencies - """ - if isscalar(frequencies): - frequencies = array(frequencies) - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies.reshape((-1, 1)) - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_long = 1j*omega*mu0/(2.*pi*beta**2*gamma**2) - - npts = int(floor(abs(radius_large-radius_small)/step_size)+1) - radii = linspace(radius_small, radius_large, npts).reshape((1, -1)) - one_array = ones(radii.shape) - - x1 = k.dot(radii)/gamma - x1sq = x1**2 - x2 = nu.dot(radii) - zlong = (coef_long.dot(length / float(npts) * one_array) * - (sp.k0(x1) / sp.i0(x1) - 1. / (x1sq * (1. / 2. + eps1.dot(one_array) * - sp.kve(1, x2) / (x2 * sp.kve(0, x2))))) - ) - - return trapz(zlong, axis=1) - - -def _zdip_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, - layer: Layer, radius_small: float, - radius_large: float, length: float, - step_size: float = 1e-3) -> ArrayLike: - """ - Function to compute the transverse dip. resistive-wall impedance for a - round taper, integrating the radius-dependent approximated formula - for a cylindrical structure (see_zdip_round_single_layer_approx above), - over the length of the taper. - :param frequencies: the frequencies (array) (in Hz) - :param gamma: relativistic mass factor - :param layer: a layer with material properties (only resistivity, - relaxation time and magnetic susceptibility are taken into account - at this stage) - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param length: the total length of the taper (in m) - :param step_size: the step size (in the radial direction) for the - integration (in m) - :return: the transverse dipolar impedance at these frequencies - """ - if isscalar(frequencies): - frequencies = array(frequencies) - beta = sqrt(1.-1./gamma**2) - omega = 2*pi*frequencies.reshape((-1,1)) - k = omega/(beta*c_light) - - rho = layer.dc_resistivity - tau = layer.resistivity_relaxation_time - mu1 = 1.+layer.magnetic_susceptibility - eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) - nu = k*sqrt(1.-beta**2*eps1*mu1) - - coef_dip = 1j*k**2*Z0/(4.*pi*beta*gamma**4) - - npts = int(floor(abs(radius_large-radius_small)/step_size)+1) - radii = linspace(radius_small,radius_large,npts).reshape((1,-1)) - one_array = ones(radii.shape) - - x1 = k.dot(radii)/gamma - x1sq = x1**2 - x2 = nu.dot(radii) - zdip = ( - coef_dip.dot(length / float(npts) * one_array) * - (sp.k1(x1) / sp.i1(x1) + 4 * beta**2 * gamma**2 / (x1sq * (2 + x2 * sp.kve(0, x2) / (mu1 * sp.kve(1, x2))))) - ) - - return trapz(zdip, axis=1) - - -def create_taper_RW_approx_component(plane: str, exponents: Tuple[int, int, int, int], - input_data: Union[FlatIW2DInput, RoundIW2DInput], - radius_small: float, radius_large: float, - step_size: float = 1e-3) -> Component: - """ - Creates a single component object modeling a round or flat taper (flatness - along the horizontal direction, change of half-gap along the vertical one) - resistive-wall impedance, using the integration of the radius-dependent - approximated formula for a cylindrical structure (see - the above functions), over the length of the taper. - :param plane: the plane the component corresponds to - :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) - :param input_data: an IW2D input object (flat or round). If the input - is of type FlatIW2DInput and symmetric, we apply to the round formula the - Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, - KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, - https://cds.cern.ch/record/248630/files/p221.pdf), - while for a single plate we use those from A. Burov and V. Danilov, - PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other - kinds of asymmetric structure will raise an error. - If the input is of type RoundIW2DInput, the structure is in principle - round but the Yokoya factors put in the input will be used. - Note that the radius or half-gaps in input_data are not used (replaced - by the scan from radius_small to radius_large, for the integration). - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param step_size: the step size (in the radial or vertical direction) - for the integration (in m) - :return: A component object - """ - gamma = input_data.relativistic_gamma - length = input_data.length - - if isinstance(input_data, FlatIW2DInput): - if len(input_data.top_layers) > 1: - raise NotImplementedError("Input data can have only one layer") - yok_long = 1. - layer = input_data.top_layers[0] - radius = input_data.top_half_gap - if input_data.top_bottom_symmetry: - yok_dipx = pi**2/24. - yok_dipy = pi**2/12. - yok_quax = -pi**2/24. - yok_quay = pi**2/24. - elif input_data.bottom_half_gap == inf: - yok_dipx = 0.25 - yok_dipy = 0.25 - yok_quax = -0.25 - yok_quay = 0.25 - else: - raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " - "hence the bottom half gap must be infinite") - elif isinstance(input_data, RoundIW2DInput): - radius = input_data.inner_layer_radius - if len(input_data.layers) > 1: - raise NotImplementedError("Input data can have only one layer") - layer = input_data.layers[0] - yok_long = input_data.yokoya_factors[0] - yok_dipx = input_data.yokoya_factors[1] - yok_dipy = input_data.yokoya_factors[2] - yok_quax = input_data.yokoya_factors[3] - yok_quay = input_data.yokoya_factors[4] - else: - raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") - - # Longitudinal impedance - if plane == 'z' and exponents == (0, 0, 0, 0): - impedance = lambda f: yok_long*_zlong_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - # Transverse impedances - elif plane == 'x' and exponents == (1, 0, 0, 0): - impedance = lambda f: yok_dipx*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'y' and exponents == (0, 1, 0, 0): - impedance = lambda f: yok_dipy*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'x' and exponents == (0, 0, 1, 0): - impedance = lambda f: yok_quax*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - elif plane == 'y' and exponents == (0, 0, 0, 1): - impedance = lambda f: yok_quay*_zdip_round_taper_RW_approx( - f, gamma, layer, radius_small, radius_large, - length, step_size=step_size) - else: - impedance = lambda f: 0 - - return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], - test_exponents=exponents[2:]) - - -def create_taper_RW_approx_element( - input_data: Union[FlatIW2DInput, RoundIW2DInput], - beta_x: float, beta_y: float, - radius_small: float, radius_large: float, step_size: float=1e-3, - component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), - name: str = "", tag: str = "", description: str = "") -> Element: - """ - Creates an element object modeling a round or flat taper (flatness - along the horizontal direction, change of half-gap along the vertical one) - resistive-wall impedance, using the integration of the radius-dependent - approximated formula for a cylindrical structure (see - the above functions), over the length of the taper. - :param input_data: an IW2D input object (flat or round). If the input - is of type FlatIW2DInput and symmetric, we apply to the round formula the - Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, - KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, - https://cds.cern.ch/record/248630/files/p221.pdf), - while for a single plate we use those from A. Burov and V. Danilov, - PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other - kinds of asymmetric structure will raise an error. - If the input is of type RoundIW2DInput, the structure is in principle round - but the Yokoya factors put in the input will be used. - Note that the radius or half-gaps in input_data are not used (replaced - by the scan from radius_small to radius_large, for the integration). - :param beta_x: The beta function in the x-plane at the position of the element - :param beta_y: The beta function in the y-plane at the position of the element - :param radius_small: the smallest radius of the taper (in m) - :param radius_large: the largest radius of the taper (in m) - :param step_size: the step size (in the radial or vertical direction) - for the integration (in m) - :param component_ids: a list of components to be computed - :param name: A user-specified name for the Element - :param tag: A string to tag the Element - :param description: A description for the Element - :return: An Element object representing the structure - """ - components = [] - length = input_data.length - for component_id in component_ids: - _, plane, exponents = component_names[component_id] - components.append(create_taper_RW_approx_component(plane=plane, exponents=exponents, input_data=input_data, - radius_small=radius_small, radius_large=radius_large, - step_size=step_size)) - - return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, - description=description) - - -def create_interpolated_impedance_component(interpolation_frequencies: ArrayLike, - impedance: Optional[Callable] = None, - wake: Optional[Callable] = None, plane: str = '', - source_exponents: Tuple[int, int] = (-1, -1), - test_exponents: Tuple[int, int] = (-1, -1), - name: str = "Unnamed Component", - f_rois: Optional[List[Tuple[float, float]]] = None, - t_rois: Optional[List[Tuple[float, float]]] = None): - """ - Creates a component in which the impedance function is evaluated directly only on few points and it is interpolated - everywhere else. This helps when the impedance function is very slow to evaluate. - :param interpolation_frequencies: the frequencies where the function is evaluated for the interpolation - :param impedance: A callable function representing the impedance function of the Component. Can be undefined if - the wake function is defined. - :param wake: A callable function representing the wake function of the Component. Can be undefined if - the impedance function is defined. - :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be specified for valid initialization - :param source_exponents: The exponents in the x and y planes experienced by the source particle. Also - referred to as 'a' and 'b'. Must be specified for valid initialization - :param test_exponents: The exponents in the x and y planes experienced by the source particle. Also - referred to as 'a' and 'b'. Must be specified for valid initialization - :param name: An optional user-specified name of the component - :param f_rois: a list of frequency regions of interest - :param t_rois: a list of time regions of interest - """ - def interpolated_impedance(x): - return np.interp(x, interpolation_frequencies, impedance(interpolation_frequencies)) - - return Component(impedance=interpolated_impedance, wake=wake, plane=plane, - source_exponents=source_exponents, test_exponents=test_exponents, - name=name, f_rois=f_rois, - t_rois=t_rois) +from xwakes.wit.utilities import * diff --git a/pywit/utils.py b/pywit/utils.py index cff25e4b..10834455 100644 --- a/pywit/utils.py +++ b/pywit/utils.py @@ -1,35 +1 @@ -from typing import Sequence, Union - -import numpy as np - - -def round_sigfigs(arr: Union[float, Sequence[float]], sigfigs: int): - """ - Rounds all the floats in an array to a the given number of - significant digits. - Warning: this routine is slow when working with very long - arrays. If the final goal is to get unique elements to the given - precision, it is better to use unique_sigfigs. - :param arr: A float or an array-like sequence of floats. - :param sigfigs: The number of significant digits (integer). - :return: The same array (or float) rounded appropriately. - """ - if np.isscalar(arr): - return np.round(arr, sigfigs - 1 - int(np.floor(np.log10(np.abs(arr))))) - else: - return np.array([round_sigfigs(value, sigfigs) for value in arr]) - - -def unique_sigfigs(arr: np.ndarray, sigfigs: int): - """ - Given an array, returns it sorted and with only the elements that - are unique after rounding to a given number of significant digits. - :param arr: An array of floats - :param sigfigs: The number of significant digits (integer) - :return: the sorted array with unique elements - """ - arr = np.unique(arr.flat) - sorted_indices = np.argsort(arr.flat) - log_diff = np.floor(np.log10(np.append(np.inf, np.diff(arr[sorted_indices])))) - log_arr = np.floor(np.log10(arr[sorted_indices])) - return arr.flat[sorted_indices[log_diff - log_arr >= -sigfigs]] +from xwakes.wit.utils import * diff --git a/setup.py b/setup.py index dcc6e115..b8952221 100644 --- a/setup.py +++ b/setup.py @@ -25,14 +25,16 @@ setup( name='xwakes', version=__version__, - description='Wakefield and impedance toolbox for Xsuite', - long_description='Wakefield and impedance toolbox for Xsuite', + description='Wake and impedance toolbox', + long_description='Toolbox to build and manipulate impedance and wake' + ' function models, usable in Xsuite, DELPHI and others', url='https://xsuite.web.cern.ch/', - author='G. Iadarola et al.', + author='M. Rognlien, L. Giacomel, D. Amorim, E. Vik, G. Iadarola ' + 'and N. Mounet', license='Apache 2.0', download_url="https://pypi.python.org/pypi/xwakes", project_urls={ - "Bug Tracker": "https://github.com/xsuite/xsuite/issues", + "Bug Tracker": "https://github.com/xsuite/xwakes/issues", "Source Code": "https://github.com/xsuite/xwakes/", }, packages=find_packages(), @@ -40,6 +42,7 @@ include_package_data=True, install_requires=[ 'numpy>=1.0', + 'scipy', 'pyyaml', ], extras_require={ diff --git a/tests/test_component.py b/tests/test_component.py index da01f651..adb1786b 100644 --- a/tests/test_component.py +++ b/tests/test_component.py @@ -1,4 +1,5 @@ -from pywit.component import Component +from pywit.component import (Component, + mix_fine_and_rough_sampling) from test_common import functions from pywit.parameters import * @@ -10,6 +11,61 @@ import numpy as np +def test_mix_fine_and_rough_sampling_no_rois(): + start = 1 + stop = 10 + rough_points = 10 + assert np.allclose( + np.geomspace(start, stop, rough_points), + mix_fine_and_rough_sampling(start=start, + stop=stop, + rough_points=rough_points, + fine_points=0, + rois=[])) + +def test_mix_fine_and_rough_sampling_with_rois_no_overlap(): + start = 1 + stop = 10 + rough_points = 10 + fine_points = 100 + # here there is no overlap between the rois and the rough points + # so all the values in the rois should be included in the total array + rois = [(5.1, 5.9), (6.1, 6.9)] + target_val = np.sort( + np.concatenate((np.geomspace(start, stop, rough_points), + np.linspace(rois[0][0], rois[0][1], fine_points), + np.linspace(rois[1][0], rois[1][1], fine_points)))) + assert np.allclose( + target_val, + mix_fine_and_rough_sampling(start=start, + stop=stop, + rough_points=rough_points, + fine_points=fine_points, + rois=rois)) + +def test_mix_fine_and_rough_sampling_with_rois_with_overlap(): + start = 1 + stop = 10 + rough_points = 10 + fine_points = 100 + rough_points_arr = np.geomspace(start, stop, rough_points) + # in the first ROI the extrema overlap with two vales in the rough_points + # up to the 7th significant digit so they should not appear in the total array + rois = [(rough_points_arr[4] + 1e-7, + rough_points_arr[5] + 1e-7), + (6.1, 6.9)] + target_val = np.sort( + np.concatenate((rough_points_arr, + np.linspace(rois[0][0], rois[0][1], fine_points)[1:-1], + np.linspace(rois[1][0], rois[1][1], fine_points)))) + assert np.allclose( + target_val, + mix_fine_and_rough_sampling(start=start, + stop=stop, + rough_points=rough_points, + fine_points=fine_points, + rois=rois)) + def simple_component_from_rois(f_rois=None, t_rois=None): return Component(impedance=lambda f: f+1j*f, wake=lambda z: z, plane='z', source_exponents=(0, 0), diff --git a/tests/test_interface.py b/tests/test_interface.py index f471884c..2c18ea14 100644 --- a/tests/test_interface.py +++ b/tests/test_interface.py @@ -15,6 +15,13 @@ import numpy as np +CURR_DIR = Path(__file__).resolve().parent + +try: + import IW2D + iw2d_installed = True +except ImportError: + iw2d_installed = False @mark.parametrize('is_impedance, plane, exponents, expected_comp_name', [[True, 'z', (0, 0, 0, 0), 'zlong'], @@ -39,15 +46,10 @@ def test_get_component_name_raise(is_impedance, plane, exponents): " the values of component_names dictionary" +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_duplicate_component_iw2d_import(): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") - with raises(AssertionError) as error_message: - import_data_iw2d(directory=Path("test/test_data/iw2d/duplicate_components").resolve(), + import_data_iw2d(directory=CURR_DIR / "test_data/iw2d/duplicate_components", common_string="WLHC_2layersup_0layersdown6.50mm") assert error_message.value.args[0] in ["The wake files 'WlongWLHC_2layersup_0layersdown6.50mm.dat' and " @@ -59,35 +61,27 @@ def test_duplicate_component_iw2d_import(): ] +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_no_matching_filename_iw2d_import(): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") + test_dir = CURR_DIR / "test_data/iw2d/valid_directory" with raises(AssertionError) as error_message: - import_data_iw2d(directory=Path("test/test_data/iw2d/valid_directory").resolve(), + import_data_iw2d(directory=test_dir, common_string="this_string_matches_no_file") expected_error_message = f"No files in " \ - f"'{Path('test/test_data/iw2d/valid_directory').resolve()}'" \ + f"'{test_dir.resolve()}'" \ f" matched the common string 'this_string_matches_no_file'." assert error_message.value.args[0] == expected_error_message +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_valid_iw2d_component_import(): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") - # Normally, the relativistic gamma would be an attribute of a required IW2DInput object, but here it has been # hard-coded instead relativstic_gamma = 479.605064966 - recipes = import_data_iw2d(directory=Path("test/test_data/iw2d/valid_directory").resolve(), + recipes = import_data_iw2d(directory=CURR_DIR / "test_data/iw2d/valid_directory", common_string="precise") for recipe in recipes: component = create_component_from_data(*recipe, relativistic_gamma=relativstic_gamma) @@ -110,7 +104,7 @@ def iw2d_input(request): return RoundIW2DInput(machine='test', length=1, relativistic_gamma=7000, calculate_wake=request.param['wake_computation'], f_params=f_params, comment='test', layers=layers_tung, inner_layer_radius=5e-2, yokoya_factors=(1, 1, 1, 1, 1), z_params=z_params) - + if request.param['chamber_type'] == 'flat': return FlatIW2DInput(machine='test', length=1, relativistic_gamma=7000, calculate_wake=request.param['wake_computation'], f_params=f_params, comment='test', @@ -133,13 +127,8 @@ def _remove_non_empty_directory(directory_path: Path): ({'chamber_type': 'round', 'wake_computation': True}, ['Zlong', 'Zxdip', 'Zydip', 'Zxquad', 'Zyquad', 'Wlong', 'Wxdip', 'Wydip', 'Wxquad', 'Wyquad']), ({'chamber_type': 'flat', 'wake_computation': True}, ['Zlong', 'Zxdip', 'Zydip', 'Zxquad', 'Zyquad', 'Zycst', 'Wlong', 'Wxdip', 'Wydip', 'Wxquad', 'Wyquad', 'Wycst'])] @pytest.mark.parametrize("iw2d_input, components_to_test", list_of_inputs_to_test, indirect=["iw2d_input"]) +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_check_already_computed(iw2d_input, components_to_test): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") - name = 'test_hash' # create the expected directories for the dummy input @@ -187,13 +176,8 @@ def test_check_already_computed(iw2d_input, components_to_test): @pytest.mark.parametrize("iw2d_input", [{'chamber_type': 'round', 'wake_computation': False}], indirect=["iw2d_input"]) +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_add_iw2d_input_to_database(iw2d_input): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") - projects_path = Path(get_iw2d_config_value('project_directory')) input_hash = sha256(iw2d_input.__str__().encode()).hexdigest() directory_level_1 = projects_path.joinpath(input_hash[0:2]) @@ -214,13 +198,8 @@ def test_add_iw2d_input_to_database(iw2d_input): os.rmdir(directory_level_1) +@pytest.mark.skipif(not iw2d_installed, reason="IW2D is not installed") def test_check_valid_working_directory(): - - try: - import IW2D - except ImportError: - pytest.skip("IW2D is not installed") - projects_path = Path(get_iw2d_config_value('project_directory')) working_directory = projects_path.joinpath(Path("a/wrong/working_directory")) diff --git a/tests/test_utilities.py b/tests/test_utilities.py index 4dba81de..c32bbe20 100644 --- a/tests/test_utilities.py +++ b/tests/test_utilities.py @@ -3,13 +3,14 @@ create_resistive_wall_single_layer_approx_component, create_resistive_wall_single_layer_approx_element, create_taper_RW_approx_component, - create_taper_RW_approx_element, - _zlong_round_single_layer_approx, - _zdip_round_single_layer_approx) + create_taper_RW_approx_element) +from xwakes.wit.utilities import (_zlong_round_single_layer_approx, + _zdip_round_single_layer_approx) from test_common import relative_error from pywit.parameters import * -from pywit.interface import (Layer, _IW2DInputBase, FlatIW2DInput, - RoundIW2DInput, Sampling, component_names) +from pywit.interface import (FlatIW2DInput, RoundIW2DInput, Sampling, + component_names) +from xwakes.wit.interface import _IW2DInputBase from pywit.materials import layer_from_json_material_library, copper_at_temperature from typing import Dict diff --git a/xwakes/.github/workflows/pr_workflow.yaml b/xwakes/.github/workflows/pr_workflow.yaml new file mode 100644 index 00000000..6e2371a7 --- /dev/null +++ b/xwakes/.github/workflows/pr_workflow.yaml @@ -0,0 +1,24 @@ +# On PR run the tests + +on: + pull_request: + branches: + - main + +jobs: + test: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v2 + - name: Set up Python 3.11 + uses: actions/setup-python@v2 + with: + python-version: 3.11 + - name: Install dependencies + run: | + python -m pip install --upgrade pip + pip install . + - name: Run tests + run: | + cd tests + pytest -v diff --git a/xwakes/__init__.py b/xwakes/__init__.py index e69de29b..ae263b13 100644 --- a/xwakes/__init__.py +++ b/xwakes/__init__.py @@ -0,0 +1,3 @@ +from ._version import __version__ + +from .init_pywit_directory import initialize_pywit_directory diff --git a/xwakes/_version.py b/xwakes/_version.py index c57bfd58..ffcc925a 100644 --- a/xwakes/_version.py +++ b/xwakes/_version.py @@ -1 +1 @@ -__version__ = '0.0.0' +__version__ = '0.0.3' diff --git a/xwakes/init_pywit_directory.py b/xwakes/init_pywit_directory.py new file mode 100644 index 00000000..1b02a8aa --- /dev/null +++ b/xwakes/init_pywit_directory.py @@ -0,0 +1,19 @@ +import pathlib +import pickle +import os + +def initialize_pywit_directory() -> None: + home_path = pathlib.Path.home() + paths = [home_path.joinpath('pywit').joinpath(ext) for ext in ('config', 'IW2D/bin', 'IW2D/projects')] + for path in paths: + os.makedirs(path, exist_ok=True) + with open(pathlib.Path(paths[0]).joinpath('iw2d_settings.yaml'), 'w') as file: + file.write(f'binary_directory: {paths[1]}\n' + f'project_directory: {paths[2]}') + + filenames = ('component', 'element', 'iw2d_inputs') + for filename in filenames: + open(paths[0].joinpath(f"{filename}.yaml"), 'w').close() + + with open(paths[2].joinpath('hashmap.pickle'), 'wb') as handle: + pickle.dump(dict(), handle, protocol=pickle.HIGHEST_PROTOCOL) diff --git a/xwakes/wit/__init__.py b/xwakes/wit/__init__.py new file mode 100644 index 00000000..94937423 --- /dev/null +++ b/xwakes/wit/__init__.py @@ -0,0 +1,13 @@ +from . import component +from . import devices +from . import element +from . import elements_group +from . import interface +from . import landau_damping +from . import materials +from . import model +from . import parameters +from . import plot +from . import sacherer_formula +from . import utilities +from . import utils \ No newline at end of file diff --git a/xwakes/wit/component.py b/xwakes/wit/component.py new file mode 100644 index 00000000..d91ea08f --- /dev/null +++ b/xwakes/wit/component.py @@ -0,0 +1,387 @@ +from __future__ import annotations + +from .parameters import * +from .utils import unique_sigfigs + +from typing import Optional, Callable, Tuple, Union, List + +import numpy as np + + +def mix_fine_and_rough_sampling(start: float, stop: float, rough_points: int, + fine_points: int, rois: List[Tuple[float, float]]): + """ + Mix a fine and rough (geometric) sampling between start and stop, + refined in the regions of interest rois. + :param start: The lowest bound of the sampling + :param stop: The highest bound of the sampling + :param rough_points: The total number of data points to be + generated for the rough grid, between start + and stop. + :param fine_points: The number of points in the fine grid of each roi. + :param rois: List of unique tuples with the lower and upper bound of + of each region of interest. + :return: An array with the sampling obtained. + """ + intervals = [np.linspace(max(i, start), min(f, stop), fine_points) + for i, f in rois + if (start <= i <= stop or start <= f <= stop)] + rough_sampling = np.geomspace(start, stop, rough_points) + # the following concatenates adds the rough points to the fine sampling and sorts + # the result. Then duplicates are removed, where two points are considered + # equal if they are within 7 significant figures of each other. + return unique_sigfigs( + np.sort(np.hstack((*intervals, rough_sampling)),kind='stable'), 7) + + +class Component: + """ + A data structure representing the impedance- and wake functions of some Component in a specified plane. + """ + + def __init__(self, impedance: Optional[Callable] = None, wake: Optional[Callable] = None, plane: str = '', + source_exponents: Tuple[int, int] = (-1, -1), test_exponents: Tuple[int, int] = (-1, -1), + name: str = "Unnamed Component", f_rois: Optional[List[Tuple[float, float]]] = None, + t_rois: Optional[List[Tuple[float, float]]] = None): + """ + The initialization function for the Component class. + :param impedance: A callable function representing the impedance function of the Component. Can be undefined if + the wake function is defined. + :param wake: A callable function representing the wake function of the Component. Can be undefined if + the impedance function is defined. + :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be specified for valid initialization + :param source_exponents: The exponents in the x and y planes experienced by the source particle. Also + referred to as 'a' and 'b'. Must be specified for valid initialization + :param test_exponents: The exponents in the x and y planes experienced by the source particle. Also + referred to as 'a' and 'b'. Must be specified for valid initialization + :param name: An optional user-specified name of the component + """ + # Enforces that either impedance or wake is defined. + assert impedance or wake, "The impedance- and wake functions cannot both be undefined." + # The impedance- and wake functions as callable objects, e.g lambda functions + self.impedance = impedance + self.wake = wake + self.name = name + + # The plane of the Component, either 'x', 'y' or 'z' + assert plane.lower() in ['x', 'y', 'z'], \ + "Cannot initialize Component object without specified plane" + self.plane = plane + + assert source_exponents != (-1, -1) and len(source_exponents) == 2, \ + "Cannot initialize Component object without specified source exponents (a, b)" + self.source_exponents = source_exponents + assert test_exponents != (-1, -1) and len(test_exponents) == 2, \ + "Cannot initialize Component object without specified test exponents (c, d)" + self.test_exponents = test_exponents + self.power_x = (source_exponents[0] + test_exponents[0] + (plane == 'x')) / 2 + self.power_y = (source_exponents[1] + test_exponents[1] + (plane == 'y')) / 2 + self.f_rois = f_rois if f_rois else [] + self.t_rois = t_rois if t_rois else [] + + def generate_wake_from_impedance(self) -> None: + """ + Uses the impedance function of the Component object to generate its wake function, using + a Fourier transform. + :return: Nothing + """ + # # If the object already has a wake function, there is no need to generate it. + # if self.wake: + # pass + # # In order to generate a wake function, we need to make sure the impedance function is defined. + # assert self.impedance, "Tried to generate wake from impedance, but impedance is not defined." + # + # raise NotImplementedError + + # Temporary solution to avoid crashes + self.wake = lambda x: 0 + + def generate_impedance_from_wake(self) -> None: + """ + Uses the wake function of the Component object to generate its impedance function, using + a Fourier transform. + :return: Nothing + """ + # # If the object already has an impedance function, there is no need to generate it. + # if self.impedance: + # pass + # # In order to generate an impedance function, we need to make sure the wake function is defined. + # assert self.wake, "Tried to generate impedance from wake, but wake is not defined." + # + # raise NotImplementedError + + # Temporary solution to avoid crashes + self.impedance = lambda x: 0 + + def is_compatible(self, other: Component) -> bool: + """ + Compares all parameters of the self-object with the argument of the function and returns True if all of their + attributes, apart from the impedance and wake functions, are identical. + i.e. Returns True if self and other are compatible for Component addition, False otherwise + :param other: Another Component + :return: True if self and other can be added together, False otherwise + """ + return all([self.source_exponents == other.source_exponents, + self.test_exponents == other.test_exponents, + self.plane == other.plane, + self.power_x == other.power_x, + self.power_y == other.power_y]) + + def __add__(self, other: Component) -> Component: + """ + Defines the addition operator for two Components + :param self: The left addend + :param other: The right addend + :return: A new Component whose impedance and wake functions are the sums + of the respective functions of the two addends. + """ + # Enforce that the two addends are in the same plane + assert self.plane == other.plane, "The two addends correspond to different planes and cannot be added.\n" \ + f"{self.plane} != {other.plane}" + + # Enforce that the two addends have the same exponent parameters + assert self.source_exponents == other.source_exponents and self.test_exponents == other.test_exponents, \ + "The two addends have different exponent parameters and cannot be added.\n" \ + f"Source: {self.source_exponents} != {other.source_exponents}\n" \ + f"Test: {self.test_exponents} != {other.test_exponents}" + + # Defines an empty array to hold the two summed functions + sums = [] + + # Iterates through the two pairs of functions: impedances, then wakes + for left, right in zip((self.impedance, self.wake), (other.impedance, other.wake)): + # If neither addend has a defined function, we will not bother to calculate that of the sum + if (not left) and (not right): + sums.append(None) + else: + # # Generates the missing function for the addend which is missing it + # if not left: + # [self.generate_impedance_from_wake, self.generate_wake_from_impedance][len(sums)]() + # elif not right: + # [other.generate_impedance_from_wake, other.generate_wake_from_impedance][len(sums)]() + # + + # TODO: Temporary fix until Fourier transform implemented + if not left: + sums.append(right) + elif not right: + sums.append(left) + else: + # Appends the sum of the functions of the two addends to the list "sums" + sums.append(lambda x, l=left, r=right: l(x) + r(x)) + + # Initializes and returns a new Component + return Component(sums[0], sums[1], self.plane, self.source_exponents, self.test_exponents, + f_rois=self.f_rois + other.f_rois, t_rois=self.t_rois + other.t_rois) + + def __radd__(self, other: Union[int, Component]) -> Component: + """ + Implements the __radd__ method for the Component class. This is only done to facilitate the syntactically + practical use of the sum() method for Components. sum(iterable) works by adding all of the elements of the + iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case + that the left addend of any addition involving a Component is not itself a Component, the resulting sum + is simply defined to be the right addend. + :param other: The left addend of an addition + :return: The sum of self and other if other is a Component, otherwise just self. + """ + # Checks if the left addend, other, is not a Component + if not isinstance(other, Component): + # In which case, the right addend is simply returned + return self + + # Otherwise, their sum is returned (by invocation of Component.__add__(self, other)) + return self + other + + def __mul__(self, scalar: complex) -> Component: + """ + Defines the behavior of multiplication of a Component by some, potentially complex, scalar + :param scalar: A scalar value to be multiplied with some Component + :return: A newly initialized Component identical to self in every way apart from the impedance- + and wake functions, which have been multiplied by the scalar. + """ + # Throws an AssertionError if scalar is not of the type complex, float or int + assert isinstance(scalar, complex) or isinstance(scalar, float) or isinstance(scalar, int) + # Initializes and returns a new Component with attributes like self, apart from the scaled functions + return Component((lambda x: scalar * self.impedance(x)) if self.impedance else None, + (lambda x: scalar * self.wake(x)) if self.wake else None, self.plane, + self.source_exponents, self.test_exponents, self.name, self.f_rois, self.t_rois) + + def __rmul__(self, scalar: complex) -> Component: + """ + Generalizes scalar multiplication of Component to be possibly from left and right. Both of these operations + are identical. + :param scalar: A scalar value to be multiplied with some Component + :return: The result of calling Component.__mul__(self, scalar): A newly initialized Component identical to self + in every way apart from the impedance- and wake functions, which have been multiplied by the scalar. + """ + # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function + return self * scalar + + def __truediv__(self, scalar: complex) -> Component: + """ + Implements the __truediv__ method for the Component class in order to produce the expected behavior when + dividing a Component by some scalar. That is, the scalar multiplication of the multiplicative inverse of + the scalar. + :param scalar: A scalar value for the Component to be divided by + :return: The result of calling Component.__mul__(self, 1 / scalar): A newly initialized Component identical to + self in every way apart from the impedance- and wake functions, which have been multiplied by (1 / scalar). + """ + # Defines the operation c / z to be equivalent to c * (1 / z) for some Component c and scalar z. + return self * (1 / scalar) + + def __str__(self) -> str: + """ + Implements the __str__ method for the Component class, providing an informative printout of the attributes + of the Component. + :return: A multi-line string containing information about the attributes of self, including which of the + impedance- and wake functions are defined. + """ + return f"{self.name} with parameters:\nPlane:\t\t\t{self.plane}\n" \ + f"Source exponents:\t{', '.join(str(i) for i in self.source_exponents)}\n" \ + f"Test exponents:\t\t{', '.join(str(i) for i in self.test_exponents)}\n" \ + f"Impedance function:\t{'DEFINED' if self.impedance else 'UNDEFINED'}\n" \ + f"Wake function:\t\t{'DEFINED' if self.wake else 'UNDEFINED'}\n" \ + f"Impedance-regions of interest: {', '.join(str(x) for x in self.f_rois)}\n" \ + f"Wake-regions of interest: {', '.join(str(x) for x in self.t_rois)}" + + def __lt__(self, other: Component) -> bool: + """ + Implements the __lt__ (less than) method for the Component class. This has no real physical interpretation, + it is just syntactically practical when iterating through the Components when adding two Elements. + (see Element.__add__) + :param other: The right hand side of a "<"-inequality + :return: True if the self-Component is "less than" the other-Component by some arbitrary, but consistent, + comparison. + """ + # The two Components are compared by their attributes + return [self.plane, self.source_exponents, self.test_exponents] < \ + [other.plane, other.source_exponents, other.test_exponents] + + def __eq__(self, other: Component) -> bool: + """ + Implements the __eq__ method for the Component class. Two Components are designated as "equal" if the following + two conditions hold: + 1. They are compatible for addition. That is, all of their attributes, apart from impedance- and wake functions, + are exactly equal. + 2. The resulting values from evaluating the impedance functions of the two components respectively for 50 + points between 1 and 10000 all need to be close within some given tolerance. The same has to hold for the wake + function. + This somewhat approximated numerical approach to the equality comparator aims to compensate for small + numerical/precision errors accumulated for two Components which have taken different "paths" to what should + analytically be identical Components. + :param other: The right hand side of the equality comparator + :return: True if the two Components have identical attributes and sufficiently close functions, False otherwise + """ + # First checks if the two Components are compatible for addition, i.e. if they have the same non-function + # attributes + if not self.is_compatible(other): + # If they do not, they are not equal + return False + + # Creates a numpy array of 50 points and verifies that the evaluations of the functions of the two components + # for all of these components are sufficiently close. If they are, True is returned, otherwise, False is + # returned. + xs = np.linspace(1, 10000, 50) + return (np.allclose(self.impedance(xs), other.impedance(xs), rtol=REL_TOL, atol=ABS_TOL) and + np.allclose(self.wake(xs), other.wake(xs), rtol=REL_TOL, atol=ABS_TOL)) + + def impedance_to_array(self, rough_points: int, start: float = MIN_FREQ, + stop: float = MAX_FREQ, + precision_factor: float = FREQ_P_FACTOR) -> Tuple[np.ndarray, np.ndarray]: + """ + Produces a frequency grid based on the f_rois attribute of the component and evaluates the component's + impedance function at these frequencies. + :param rough_points: The total number of data points to be + generated for the rough grid, between start + and stop. + :param start: The lowest frequency in the desired frequency grid + :param stop: The highest frequency in the desired frequency grid + :param precision_factor: A number indicating the ratio of points + which should be placed within the regions of interest. + If =0, the frequency grid will ignore the intervals in f_rois. + If =1, the points will be distributed 1:1 between the + rough grid and the fine grid on each roi. In general, =n + means that there will be n times more points in the fine + grid of each roi, than in the rough grid. + :return: A tuple of two numpy arrays with same shape, giving + the frequency grid and impedances respectively + """ + if len(self.f_rois) == 0: + xs = np.geomspace(start, stop, rough_points) + return xs, self.impedance(xs) + + # eliminate duplicates + f_rois_no_dup = set(self.f_rois) + + fine_points_per_roi = int(round(rough_points * precision_factor)) + + xs = mix_fine_and_rough_sampling(start, stop, rough_points, + fine_points_per_roi, + list(f_rois_no_dup)) + + return xs, self.impedance(xs) + + def wake_to_array(self, rough_points: int, start: float = MIN_TIME, + stop: float = MAX_TIME, + precision_factor: float = TIME_P_FACTOR) -> Tuple[np.ndarray, np.ndarray]: + """ + Produces a time grid based on the t_rois attribute of the component and evaluates the component's + wake function at these time points. + :param rough_points: The total number of data points to be + generated for the rough grid, between start + and stop. + :param start: The lowest time in the desired time grid + :param stop: The highest time in the desired time grid + :param precision_factor: A number indicating the ratio of points + which should be placed within the regions of interest. + If =0, the frequency grid will ignore the intervals in t_rois. + If =1, the points will be distributed 1:1 between the + rough grid and the fine grid on each roi. In general, =n + means that there will be n times more points in the fine + grid of each roi, than in the rough grid. + :return: A tuple of two numpy arrays with same shape, giving the + time grid and wakes respectively + """ + if len(self.t_rois) == 0: + xs = np.geomspace(start, stop, rough_points) + return xs, self.wake(xs) + + # eliminate duplicates + t_rois_no_dup = set(self.t_rois) + + fine_points_per_roi = int(round(rough_points * precision_factor)) + + xs = mix_fine_and_rough_sampling(start, stop, rough_points, + fine_points_per_roi, + list(t_rois_no_dup)) + + return xs, self.wake(xs) + + def discretize(self, freq_points: int, time_points: int, freq_start: float = MIN_FREQ, freq_stop: float = MAX_FREQ, + time_start: float = MIN_TIME, time_stop: float = MAX_TIME, + freq_precision_factor: float = FREQ_P_FACTOR, + time_precision_factor: float = TIME_P_FACTOR) -> Tuple[ + Tuple[np.ndarray, np.ndarray], Tuple[np.ndarray, np.ndarray]]: + """ + Combines the two discretization-functions in order to fully discretize the wake and impedance of the object + as specified by a number of parameters. + :param freq_points: The total number of frequency/impedance points + :param time_points: The total number of time/wake points + :param freq_start: The lowest frequency in the frequency grid + :param freq_stop: The highest frequency in the frequency grid + :param time_start: The lowest time in the time grid + :param time_stop: The highest time in the time grid + :param freq_precision_factor: The ratio of points in the fine frequency grid + (#points in frequency ROIs / rough frequency points) + :param time_precision_factor: The ratio of points in the fine time grid + (#points in time ROIs / rough time points) + :return: A tuple of two tuples, containing the results of impedance_to_array and wake_to_array respectively + """ + return (self.impedance_to_array(freq_points, freq_start, freq_stop, freq_precision_factor), + self.wake_to_array(time_points, time_start, time_stop, time_precision_factor)) + + def get_shorthand_type(self) -> str: + """ + :return: A 5-character string indicating the plane as well as source- and test exponents of the component + """ + return self.plane + "".join(str(x) for x in (self.source_exponents + self.test_exponents)) diff --git a/xwakes/wit/devices.py b/xwakes/wit/devices.py new file mode 100644 index 00000000..8d30bc8e --- /dev/null +++ b/xwakes/wit/devices.py @@ -0,0 +1,312 @@ +from typing import Tuple,Sequence +from .component import Component +from .element import Element +from .utilities import create_resonator_component +from .interface import component_names, get_component_name + +import numpy as np +from scipy import integrate + +from scipy.constants import c as c_light, mu_0 +from scipy.special import erfc, zeta + +free_space_impedance = mu_0 * c_light + + +def create_tesla_cavity_component(plane: str, exponents: Tuple[int, int, int, int], a: float, g: float, + period_length: float) -> Component: + """ + Creates a single component object modeling a periodic accelerating stucture. + Follow K. Bane formalism developped in SLAC-PUB-9663, "Short-range dipole wakefields + in accelerating structures for the NLC". + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param a: accelerating structure iris gap in m + :param g: individual cell gap in m + :param period_length: period length in m + :return: A component object of a periodic accelerating structure + """ + + # Material properties required for the skin depth computation are derived from the input Layer attributes + # material_resistivity = layer.dc_resistivity + # material_relative_permeability = layer.magnetic_susceptibility + # material_permeability = material_relative_permeability * scipy.constants.mu_0 + + # Create the skin depth as a function of frequency and layer properties + # delta_skin = lambda freq: (material_resistivity/ (2*pi*abs(freq) * material_permeability)) ** (1/2) + + gamma = g / period_length + alpha1 = 0.4648 + alpha = 1 - alpha1 * np.sqrt(gamma) - (1 - 2*alpha1) * gamma + + s00 = g / 8 * (a / (alpha * period_length)) ** 2 + + # Longitudinal impedance and wake + if plane == 'z' and exponents == (0, 0, 0, 0): + def longitudinal_impedance_tesla_cavity(freq): + return (1j * free_space_impedance / (np.pi * (2 * np.pi * freq / c_light) * a ** 2) * + (1 + (1 + 1j) * alpha * period_length / a * np.sqrt(np.pi / ((2 * np.pi * freq / c_light) * g))) ** (-1)) + + def longitudinal_wake_tesla_cavity(time): + return ((free_space_impedance * c_light / (np.pi * a ** 2)) * np.heaviside(time, 0) * + np.exp(-np.pi * c_light * time / (4 * s00)) * erfc(np.sqrt(np.pi * c_light * time / (4 * s00)))) + + impedance = longitudinal_impedance_tesla_cavity + wake = longitudinal_wake_tesla_cavity + # Transverse dipolar impedance and wake + elif (plane == 'x' and exponents == (1, 0, 0, 0)) or (plane == 'y' and exponents == (0, 1, 0, 0)): + + def transverse_dipolar_impedance_tesla_cavity(freq): + return 2 / ((2 * np.pi * freq / c_light) * a ** 2) * 1j * free_space_impedance / ( + np.pi * (2 * np.pi * freq / c_light) * a ** 2) * ( + 1 + (1 + 1j) * alpha * period_length / a * np.sqrt(np.pi / ((2 * np.pi * freq / c_light) * g))) ** (-1) + + def transverse_dipolar_wake_tesla_cavity(time): + return ((4 * free_space_impedance * c_light * s00) / (np.pi * a ** 4) * np.heaviside(time, 0) * + (1 - (1 + np.sqrt(c_light * time / s00)) * np.exp(-np.sqrt(c_light * time / s00)))) + + impedance = transverse_dipolar_impedance_tesla_cavity + wake = transverse_dipolar_wake_tesla_cavity + else: + print("Warning: tesla cavity impedance not implemented for component {}{}. Set to zero".format(plane, + exponents)) + + def zero_function(x): + return np.zeros_like(x) + + impedance = zero_function + wake = zero_function + + return Component(impedance=np.vectorize(lambda f: impedance(f)), + wake=np.vectorize(lambda t: wake(t)), + plane=plane, source_exponents=exponents[:2], test_exponents=exponents[2:]) + + +def _integral_stupakov(half_gap_small: float, half_gap_big: float, + half_width: float, g_index: int, g_power: int, + approximate_integrals: bool = False): + """ + Computes the Stupakov's integral for a rectangular linear taper. + Note that (g')^2 has been taken out of the integral. See Phys. Rev. STAB 10, 094401 (2007) + :param half_gap_small: the small half-gap of the taper + :param half_gap_big: the large half-gap of the taper + :param half_width: the half-width of the taper + :param g_index: the indice of the G function to use (0 is for G0=F in Stupakov's paper) + :param g_power: is the power to which 1/g is taken + :param approximate_integrals: use approximated formulas to compute the integrals. + It can be used if one assumes small half_gap_big/half_width ratio << 1 + """ + def _integrand_stupakov(g): + """ + Computes the integrand for the Stupakov integral + :param g: the half-gap of the taper + """ + x = g/half_width # half-gap over half-width ratio + + def g_addend(m): + """ + Computes the m-th addend of series defining F (when g_index=0) or G_[g_index] function from Stupakov's formulas for a + rectangular linear taper at a given x=half-gap/half-width. m can be an array, in which case the function returns an array. + See Phys. Rev. STAB 10, 094401 (2007) + :param m: the index of the addend (it can be an array) + """ + + if g_index == 0: + val = (2 * m + 1) * np.pi * x / 2 + res = ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * ( + 2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2) / (2 * m + 1) + + return res + + elif g_index == 1: + val = (2 * m + 1) * np.pi * x / 2 + res = x ** 3 * (2 * m + 1) * (4 * np.exp(-2 * val)) * (1 + np.exp(-2 * val)) / ((1 - np.exp(-2 * val)) ** 3) + return res + + elif g_index == 2: + val = (2 * m + 1) * np.pi * x / 2 + res = x ** 2 * (2 * m + 1) * ( + ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * (2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2)) + return res + + elif g_index == 3: + val = m * np.pi * x + res = x ** 2 * (2 * m) * ((1 - np.exp(-2 * val)) / (1 + np.exp(-2 * val)) * ( + 2 * (np.exp(-val)) / (1 + np.exp(-2 * val))) ** 2) + return res + + else: + raise ValueError("Pb in g_addend for Stupakov's formulas: g_index not 0, 1, 2 or 3 !") + + # Computes F (when g_index=0) or G_[g_index] function from Stupakov's formulas. + g_stupakov = 0. + old_g_stupakov = 1. + eps = 1e-5 # relative precision of the summation + incr = 10 # increment for sum computation + m = np.arange(incr) + m_limit = 1e6 + + while (abs(g_stupakov - old_g_stupakov) > eps * abs(g_stupakov)) and (m[-1] < m_limit): + g_array = g_addend(m) + old_g_stupakov = g_stupakov + g_stupakov += np.sum(g_array) + m += incr + + if m[-1] >= m_limit: + print("Warning: maximum number of elements reached in g_stupakov!", m[-1], x, ", err=", + abs((g_stupakov - old_g_stupakov) / g_stupakov)) + + return g_stupakov / (g ** g_power) + + if approximate_integrals: + if g_index == 0 and g_power == 0: + i = 7. * zeta(3, 1) / (2. * np.pi ** 2) * ( + half_gap_big - half_gap_small) # (zeta(3.,1.) is Riemann zeta function at x=3) + + elif g_index == 1 and g_power == 3: + i = (1. / (half_gap_small ** 2) - 1. / (half_gap_big ** 2)) / ( + 2. * np.pi) # approx. integral + + elif g_index in [2,3] and g_power == 2: + i = (1. / half_gap_small - 1. / half_gap_big) / (np.pi ** 2) + + else: + raise ValueError("Wrong values of g_index and g_power in _integral_stupakov") + + else: + # computes numerically the integral instead of using its approximation + i, err = integrate.quadrature(_integrand_stupakov, half_gap_small, half_gap_big, + tol=1.e-3, maxiter=200, vec_func=False) + + return i + + +def shunt_impedance_flat_taper_stupakov_formula(half_gap_small: float, half_gap_big: float, taper_slope: float, + half_width: float, cutoff_frequency: float = None, + component_id: str = None, approximate_integrals: bool = False) -> float: + """ + Computes the shunt impedance in Ohm(/m if not longitudinal) of one single rectangular linear taper using Stupakov's + formulae (Phys. Rev. STAB 10, 094401 - 2007), multiplied by Z0*c/(4*pi) to convert to SI units. + Taper is in the vertical plane (i.e. the change of half-gap is in vertical, but the width is along the horizontal plane + - this is typical of a taper going towards a vertical collimator with horizontal jaws). + We use here half values for geometrical parameters (half-gap and half-width) whereas Stupakov's paper is expressed + with full values. This does not make any difference except for an additional factor 4 here for longitudinal impedance. + + The formula is valid under the conditions of low frequency and length of taper much larger than its transverse + dimensions. + + Note: one gets zero if component_id is not in [zlong, zxdip, zydip, zxqua, zyqua]. + + :param half_gap_small: small vertical half-gap + :param half_gap_big: large vertical half-gap + :param taper_slope: the slope of the taper + :param half_width: half width of the taper (constant) + :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) + :param component_id: a component name for which one computes the R/Q + (ex: zlong, zydip, zxqua, etc.) + :param approximate_integrals: use approximated formulas to compute the integrals. It can be used if one assumes + small half_gap_big/half_width ratio (<<1) + :return: the shunt impedance of the component. + """ + z_0 = mu_0 * c_light # free space impedance + + if cutoff_frequency is None and component_id == 'zlong': + raise ValueError("cutoff_frequency must be specified when component_id is 'zlong'") + + if component_id == 'zlong': + g_index = 0 + g_power = 0 + cst = 4. * mu_0 * cutoff_frequency / 2. # factor 4 due to use of half-gaps here + + elif component_id == 'zydip': + g_index = 1 + g_power = 3 + cst = z_0 * half_width * np.pi / 4. + + elif component_id == 'zxqua': + g_index = 2 + g_power = 2 + cst = -z_0 * np.pi / 4. + + elif component_id == 'zxdip': + g_index = 3 + g_power = 2 + cst = z_0 * np.pi / 4. + + elif component_id == 'zyqua': + g_index = 2 + g_power = 2 + cst = z_0 * np.pi / 4. + else: + # mock-up values + g_power = 0 + g_index = 0 + cst = 0 + + # computes the integral + i = _integral_stupakov(half_gap_small, half_gap_big, half_width, g_index, g_power, + approximate_integrals = approximate_integrals) + i *= taper_slope # put back g' factor that was dropped + + return cst * i + + +def create_flat_taper_stupakov_formula_component(half_gap_small: float, half_gap_big: float, taper_slope: float, + half_width: float, plane: str, exponents: Tuple[int, int, int, int], + cutoff_frequency: float = None) -> Component: + """ + Creates a component using the flat taper Stupakov formula + :param half_gap_small: small vertical half-gap + :param half_gap_big: large vertical half-gap + :param taper_slope: the slope of the taper + :param half_width: half width of the taper (constant) + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) + :return: A component object of a flat taper + """ + component_id = get_component_name(True, plane, exponents) + + r_shunt = shunt_impedance_flat_taper_stupakov_formula(half_gap_small=half_gap_small, half_gap_big=half_gap_big, + taper_slope=taper_slope, half_width=half_width, + cutoff_frequency=cutoff_frequency, component_id=component_id) + + return create_resonator_component(plane=plane, exponents=exponents, r=r_shunt, q=1, f_r=cutoff_frequency) + + +def create_flat_taper_stupakov_formula_element(half_gap_small: float, half_gap_big: float, taper_slope: float, + half_width: float, beta_x: float, beta_y: float, + cutoff_frequency: float = None, + component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), + name: str = "Flat taper", tag: str = "", + description: str = "") -> Element: + """ + Creates an element using the flat taper Stupakov formula + :param half_gap_small: small vertical half-gap + :param half_gap_big: large vertical half-gap + :param taper_slope: the slope of the taper + :param half_width: half width of the taper (constant) + :param beta_x: The beta function in the x-plane at the position of the taper + :param beta_y: The beta function in the y-plane at the position of the taper + :param cutoff_frequency: the cutoff frequency (used only for the longitudinal component) + :param component_ids: a list of components to be computed + :param name: A user-specified name for the Element + :param tag: A string to tag the Element + :param description: A description for the Element + :return: An Element object representing the flat taper + """ + length = (half_gap_big - half_gap_small) / taper_slope + + components = [] + for component_id in component_ids: + _, plane, exponents = component_names[component_id] + components.append(create_flat_taper_stupakov_formula_component(half_gap_small=half_gap_small, + half_gap_big=half_gap_big, + taper_slope=taper_slope, + half_width=half_width, + plane=plane, exponents=exponents, + cutoff_frequency=cutoff_frequency, + )) + + return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, + description=description) diff --git a/xwakes/wit/element.py b/xwakes/wit/element.py new file mode 100644 index 00000000..f35309fc --- /dev/null +++ b/xwakes/wit/element.py @@ -0,0 +1,310 @@ +from __future__ import annotations + +from .component import Component, Union + +from typing import List +from collections import defaultdict + +from scipy.special import comb +import numpy as np +import copy + +class Element: + def __init__(self, length: float = 0, beta_x: float = 0, beta_y: float = 0, + components: List[Component] = None, name: str = "Unnamed Element", + tag: str = "", description: str = ""): + """ + The initialization function of the Element class. + :param length: The length of the element, must be specified for valid initialization + :param beta_x: The size of the beta function in the x-plane at the position of the Element, must be specified + for valid initialization + :param beta_y: The size of the beta function in the y-plane at the position of the Element, must be specified + for valid initialization + :param components: A list of the Components corresponding to the Element being initialized. If the list contains + multiple Components with the same values for (plane, source_exponents, test_exponents), and error is thrown. If + the list is not specified, the Element is initialized with an empty components-list. + :param name: A user-specified name of the Element + :param tag: A string corresponding to a specific component + """ + assert length > 0, "The element must have a specified valid length" + assert beta_x > 0 and beta_y > 0, "The element must have valid specified beta_x and beta_y values" + self.length = length + self.beta_x = beta_x + self.beta_y = beta_y + self.name = name + self.tag = tag + self.description = description + if components: + comp_dict = defaultdict(int) + for c in components: + comp_dict[(c.plane, c.source_exponents, c.test_exponents)] += c + components = comp_dict.values() + self.components = sorted(components, key=lambda x: (x.plane, x.source_exponents, x.test_exponents)) + else: + self.components: List[Component] = [] + + def rotated(self, theta: float, rotate_beta: bool = False) -> Element: + """ + Returns a copy if the self-Element which has been rotated counterclockwise in the transverse plane by an + angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y values of the new Element + are rotated correspondingly. + :param theta: The angle, in radians, the Element is to be rotated by + :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new Element should also + be rotated by theta. Enabled by default. + :return: A newly initialized copy of the self-element which has been rotated counterclockwise in the + transverse plane by an angle theta. + """ + # Confines theta to the interval [0, 2 * pi) + theta %= 2 * np.pi + # Precalculates the sine and cosine of theta to save time + costheta, sintheta = np.cos(theta), np.sin(theta) + # Creates a dictionary where the keys are Component attributes and the values are Components with those + # corresponding attributes. That way, Components created by the rotation can always be added to other compatible + # Components. + rotated_components = defaultdict(int) + + # Iterates through all the components of the self-Element + for cmp in self.components: + # Defines the x- and y-coefficients depending on what plane the Component is in + if cmp.plane == 'x': + coefx, coefy = costheta, sintheta + else: + coefx, coefy = -sintheta, costheta + # For compactness, a, b, c and d are used to refer to the source_exponents and the test_exponents + # of the component + a, b, c, d = cmp.source_exponents + cmp.test_exponents + for i in range(a + 1): + for j in range(b + 1): + for k in range(c + 1): + for l in range(d + 1): + # Calculates new a, b, c and d values for the new Component + new_a, new_b, new_c, new_d = i + j, a - i + b - j, k + l, c - k + d - l + # Product of binomial coefficients + binprod = int(comb(a, i, exact=True) * comb(b, j, exact=True) * + comb(c, k, exact=True) * comb(d, l, exact=True)) + # Multiply by power of cos and sin + coef = ((-1) ** (j + l)) * binprod * (costheta ** (i + b - j + k + d - l)) * \ + (sintheta ** (a - i + j + c - k + l)) + # Depending on if the component is in the longitudinal or transverse plane, one or two + # components are created and added to the element at the correct key in the + # rotated_components dictionary + if cmp.plane == 'z': + # If the component is scaled by less than 10 ** -6 we assume that it is zero + if abs(coef) > 1e-6: + rotated_components['z', new_a, new_b, new_c, new_d] += \ + coef * Component(impedance=cmp.impedance, wake=cmp.wake, plane='z', + source_exponents=(new_a, new_b), + test_exponents=(new_c, new_d)) + else: + if abs(coefx * coef) > 1e-6: + rotated_components['x', new_a, new_b, new_c, new_d] += \ + (coefx * coef) * Component(impedance=cmp.impedance, wake=cmp.wake, plane='x', + source_exponents=(new_a, new_b), + test_exponents=(new_c, new_d)) + if abs(coefy * coef) > 1e-6: + rotated_components['y', new_a, new_b, new_c, new_d] += \ + (coefy * coef) * Component(impedance=cmp.impedance, wake=cmp.wake, plane='y', + source_exponents=(new_a, new_b), + test_exponents=(new_c, new_d)) + + # New beta_x and beta_y values are defined if the rotate_beta flag is active + new_beta_x = ((costheta * np.sqrt(self.beta_x) - + sintheta * np.sqrt(self.beta_y)) ** 2) if rotate_beta else self.beta_x + new_beta_y = ((sintheta * np.sqrt(self.beta_x) + + costheta * np.sqrt(self.beta_y)) ** 2) if rotate_beta else self.beta_y + + # Initializes and returns a new element with parameters calculated above. Its Component list is a list of + # the values in the rotated_components dictionary, sorted by the key which is used consistently for + # Component comparisons + return Element(self.length, new_beta_x, new_beta_y, + sorted(list(rotated_components.values()), + key=lambda x: (x.plane, x.source_exponents, x.test_exponents)), + self.name, self.tag, self.description) + + def is_compatible(self, other: Element, verbose: bool = False) -> bool: + """ + Compares all non-components parameters of the two Elements and returns False if they are not all equal within + some tolerance. The Component lists of the two Elements also need to be of equal length in order for the + function to return True. + :param other: An element for the self-Element to be compared against + :param verbose: A flag which can be activated to give feedback on the values of the parameters which caused + the function to return False + :return: True if all non-components parameters are equal within some tolerance and the Components lists of the + two Elements are of the same length. False otherwise. + """ + if verbose: + if abs(self.length - other.length) > 1e-8: + print(f"Different lengths: {self.length} != {other.length}") + return False + if abs(self.beta_x - other.beta_x) > 1e-8: + print(f"Different beta_x: {self.beta_x} != {other.beta_x}") + return False + if abs(self.beta_y - other.beta_y) > 1e-8: + print(f"Different beta_y: {self.beta_y} != {other.beta_y}") + return False + if len(self.components) != len(other.components): + print(f"Different number of components: {len(self.components)} != {len(other.components)}") + return False + return True + else: + return all([abs(self.length - other.length) < 1e-8, + abs(self.beta_x - other.beta_x) < 1e-8, + abs(self.beta_y - other.beta_y) < 1e-8, + len(self.components) == len(other.components)]) + + def changed_betas(self, new_beta_x: float, new_beta_y: float) -> Element: + element_copy = copy.deepcopy(self) + x_ratio = self.beta_x / new_beta_x + y_ratio = self.beta_y / new_beta_y + element_copy.components = [((x_ratio ** c.power_x) * (y_ratio ** c.power_y)) * c for c in self.components] + element_copy.beta_x = new_beta_x + element_copy.beta_y = new_beta_y + return element_copy + + def __add__(self, other: Element) -> Element: + """ + Defines the addition operator for two objects of the class Element + :param self: The left addend + :param other: The right addend + :return: A new object of the class Element which represents the sum of the two addends. Its length is simply + the sum of the length of the addends, its beta functions are a weighted, by length, sum of the beta functions + of the addends, and its list of components include all of the components of the addends, added together where + possible. + """ + # Defines new attributes based on attributes of the addends + new_length = self.length + other.length + new_beta_x = (self.length * self.beta_x + other.length * other.beta_x) / new_length + new_beta_y = (self.length * self.beta_y + other.length * other.beta_y) / new_length + new_components = [] + + # Pre-calculates some ratios which will be used in multiple calculations + ratios = (self.beta_x / new_beta_x, self.beta_y / new_beta_y, + other.beta_x / new_beta_x, other.beta_y / new_beta_y) + + # i and j represent indices in the component lists of the left and right element respectively + i, j = 0, 0 + # This while loop iterates through pairs of components in the component-lists of the elements until + # one of the lists is exhausted. + while i < len(self.components) and j < len(other.components): + # If the two components are compatible for addition, they are weighted according to the beta functions + if self.components[i].is_compatible(other.components[j]): + comp1 = self.components[i] + power_x = comp1.power_x + power_y = comp1.power_y + left_coefficient = (ratios[0] ** power_x) * (ratios[1] ** power_y) + right_coefficient = (ratios[2] ** power_x) * (ratios[3] ** power_y) + new_components.append(left_coefficient * comp1 + right_coefficient * other.components[j]) + i += 1 + j += 1 + # If the left component is "less than" the right component, by some arbitrary pre-defined measure, we know + # that there cannot be a component in the right element compatible with our left component. Therefore, we + # simply add the left component to the list of new components + elif self.components[i] < other.components[j]: + left_coefficient = (ratios[0] ** self.components[i].power_x) * (ratios[1] ** self.components[i].power_y) + new_components.append(left_coefficient * self.components[i]) + i += 1 + else: + right_coefficient = ((ratios[2] ** other.components[j].power_x) * + (ratios[3] ** other.components[j].power_y)) + new_components.append(right_coefficient * other.components[j]) + j += 1 + + # When the while-loop above exits, there could still be unprocessed components remaining in either the + # component list of either the left- or the right element, but not both. We simply append any remaining + # components to our new_components + if i != len(self.components): + for c in self.components[i:]: + left_coefficient = (ratios[0] ** c.power_x) * (ratios[1] ** c.power_y) + new_components.append(left_coefficient * c) + elif j != len(other.components): + for c in other.components[j:]: + right_coefficient = (ratios[2] ** c.power_x) * (ratios[3] ** c.power_y) + new_components.append(right_coefficient * c) + + # Creates and returns a new element which represents the sum of the two added elements. + return Element(new_length, new_beta_x, new_beta_y, new_components) + + def __radd__(self, other: Union[int, Element]) -> Element: + """ + Implements the __rad__ method for the Element class. This is only done to facilitate the syntactically + practical use of the sum() method for Elements. sum(iterable) works by adding all of the elements of the + iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case + that the left addend of any addition involving an Element is not itself an Element, the resulting sum + is simply defined to be the right addend. + :param other: The left addend of an addition + :return: The sum of self and other if other is an Element, otherwise just self. + """ + # Checks if the left addend, other, is not an Element + if not isinstance(other, Element): + # In which case, the right addend is simply returned + return self + + # Otherwise, their sum is returned (by invocation of Component.__add__(self, other)) + return self + other + + def __mul__(self, scalar: float) -> Element: + """ + Implements the __mul__ method for the Element class. Defines the behavior of multiplication of an Element by + some scalar. + :param scalar: A scalar value to be multiplied with some Element (cannot be complex) + :return: A newly initialized Element which has the same beta_x and beta_y values as the self-Element, but has + its length and all of the Components in its Component list multiplied by the scalar + """ + # Multiplies the length of the self-element by the scalar + new_length = self.length * scalar + # Creates a new list of Components which is a copy of the Component list of the self-element except every + # Component is multiplied by the scalar + new_components = [c * scalar for c in self.components] + # Initializes and returns a new Element with the arguments defined above + return Element(new_length, self.beta_x, self.beta_y, new_components, self.name, self.tag, self.description) + + def __rmul__(self, scalar: float) -> Element: + """ + Generalizes scalar multiplication of Element to be possibly from left and right. Both of these operations + are identical. + :param scalar: A scalar value to be multiplied with some Element + :return: The result of calling Element.__mul__(self, scalar): A newly initialized Element which has the same + beta_x and beta_y values as the self-Element, but has its length and all of the Components in its Component list + multiplied by the scalar. + """ + # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function + return self * scalar + + def __eq__(self, other: Element) -> bool: + """ + Implements the __eq__ method for the Element class. Two Elements are designated as "equal" if the following + two conditions hold: + 1. They have equal attributes within some tolerance and the length of the components-lists need to be identical. + 2. For every pair of components in the components-lists of the two elements, the two Components need to + evaluate as equal. + This somewhat approximated empirical approach to the equality comparator aims to compensate for small + numerical/precision errors accumulated for two Elements which have taken different "paths" to what should + analytically be identical Elements. + Note that this comparator requires the evaluation of every pair of wake- and impedance functions in some + number of points for every component in each of the elements. As such, it can be, though not extremely, + somewhat computationally intensive, and should not be used excessively. + :param other: The right hand side of the equality comparator + :return: True if the two Elements have sufficiently close attributes and every pair of components evaluate as + equal by the __eq__ method for the Component class. + """ + # Verifies that the two elements have sufficiently close attributes and components-lists of the same length + if not self.is_compatible(other): + return False + + # Returns true if every pair of components in the two lists evaluate as true by the __eq__ method defined in + # the Component class + return all(c1 == c2 for c1, c2 in zip(self.components, other.components)) + + def __str__(self): + return f"{self.name} with parameters:\n" \ + f"Length:\t\t{self.length}\n" \ + f"Beta_x:\t\t{self.beta_x}\n" \ + f"Beta_y:\t\t{self.beta_y}\n" \ + f"#components:\t{len(self.components)}" + + def get_component(self, type_string: str): + for comp in self.components: + if comp.get_shorthand_type() == type_string: + return comp + + raise KeyError(f"'{self.name}' has no component of the type '{type_string}'.") diff --git a/xwakes/wit/elements_group.py b/xwakes/wit/elements_group.py new file mode 100644 index 00000000..872e340b --- /dev/null +++ b/xwakes/wit/elements_group.py @@ -0,0 +1,178 @@ +from __future__ import annotations + +import numpy as np + +from .element import Element +from .component import Component, Union + +from typing import List + + +class ElementsGroup(Element): + """ + A class used to store many elements of the same kind (i.e. Collimators, Roman Pots, Broadband resonators...). + Each of these groups require different handling of the input files and impedance computations (for some we use IW2D, + while for others we simply read the wake from a file), therefore this should be used as a base class from which + specific classes are derived. + """ + def __init__(self, elements_list: List[Element], name: str = "Unnamed Element", tag: str = "", + description: str = ""): + """ + The initialization function of the ElementsGroup class. + :param elements_list: The list of elements in this group + :param name: A user-specified name of the group + :param tag: An optional keyword that can be used to help grouping different ElementGroup's + :param description: A user-specified description of the group + """ + if len(elements_list) == 0: + raise ValueError('Elements_list cannot be empty') + + self.elements_list = elements_list + + sum_element = sum(elements_list) + + # Initialize the components by summing all components with the same values of (plane, source_exponents, + # test_exponents) + super().__init__(components=sum_element.components, name=name, tag=tag, description=description, + length=sum_element.length, beta_x=sum_element.beta_x, beta_y=sum_element.beta_y) + + def __add__(self, other: Union[Element, ElementsGroup]) -> ElementsGroup: + """ + Defines the addition operator for two objects of the class ElementsGroup or for an ElementsGroup and an Element. + :param self: The left addend + :param other: The right addend + :return: A new object of the class Element which represents the sum of the two addends. For two ElementsGroups + it is an Element given by the sum of all the elements in the two groups, while for an ElementsGroup and an + Element it is the Element given by the sum of the Element and all the elements in the ElementsGroup. + """ + if isinstance(other, Element): + elements_list = self.elements_list + [other] + elif isinstance(other, ElementsGroup): + elements_list = self.elements_list + other.elements_list + else: + raise TypeError("An ElementsGroup can only be summed with an ElementGroup or an Element") + + return ElementsGroup(elements_list, name=self.name, tag=self.tag, description=self.description) + + def __radd__(self, other: Union[int, Element, ElementsGroup]) -> ElementsGroup: + """ + Implements the __rad__ method for the ElementsGroup class. This is only done to facilitate the syntactically + practical use of the sum() method for ElementsGroup. sum(iterable) works by adding all of the elements of the + iterable to 0 sequentially. Thus, the behavior of the initial 0 + iterable[0] needs to be defined. In the case + that the left addend of any addition involving an Element is not itself an Element, the resulting sum + is simply defined to be the right addend. + :param other: The left addend of an addition + :return: The sum of self and other if other is an ElementGroup or an Element, otherwise just self. + """ + if type(other) == int and other != 0: + raise ValueError("ElementsGroup right addition can only be performed with a zero addend and it is only " + "implemented to enable the sum on a list of ElementsGroup's") + # Checks if the left addend, other, is not an Element + if not (isinstance(other, ElementsGroup) or isinstance(other, Element)): + # In which case, the right addend is simply returned + return self + + # Otherwise, their sum is returned (by invocation of ElementsGroup.__add__(self, other)) + return self + other + + def __mul__(self, scalar: float) -> ElementsGroup: + """ + Implements the __mul__ method for the ElementsGroup class. Defines the behavior of multiplication of an Element + by some scalar. + :param scalar: A scalar value to be multiplied with some Element (cannot be complex) + :return: A newly initialized GroupElement in which every element is multiplied by the scalar. + """ + mult_elements = [] + for element in self.elements_list: + mult_elements.append(element * scalar) + + return ElementsGroup(mult_elements, name=self.name, tag=self.tag, description=self.description) + + def __rmul__(self, scalar: float) -> ElementsGroup: + """ + Generalizes scalar multiplication of ElementsGroup to be possibly from left and right. Both of these operations + are identical. + :param scalar: A scalar value to be multiplied with some ElementsGroup + :return: The result of calling ElementsGroup.__mul__(self, scalar): A newly initialized GroupElement in which + every element is multiplied by the scalar. + """ + # Simply swaps the places of scalar and self in order to invoke the previously defined __mul__ function + return self * scalar + + def __eq__(self, other: ElementsGroup) -> bool: + """ + Implements the __eq__ method for the Element class. Two Elements are designated as "equal" if corresponding + elements in the elements lists are equal and if all the other non-components parameters are equal up to a small + tolerance + """ + if len(self.elements_list) != len(other.elements_list): + return False + + # Verifies that the two elements have sufficiently close attributes and components-lists of the same length + if not self.is_compatible(other): + return False + + return all(e1 == e2 for e1, e2 in zip(self.elements_list, other.elements_list)) + + def __str__(self): + string = f"{self.name} ElementsGroup composed of the following elements\n" + for element in self.elements_list: + string += str(element) + string += "============\n" + return string + + def rotated_element(self, name: str, theta: float, rotate_beta: bool = False) -> ElementsGroup: + """ + Returns a copy if the self-ElementsGroup in which a user-specified element has been rotated counterclockwise + in the transverse plane by an angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y + values of the element in the new ElementsGroup are rotated correspondingly. + :param name: the name of the Element to be rotated + :param theta: The angle, in radians, the ElementsGroup is to be rotated by + :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new ElementsGroup + should also be rotated by theta. Enabled by default. + :return: A newly initialized copy of the self-ElementsGroup which has been rotated counterclockwise in the + transverse plane by an angle theta. + """ + rotated_elements = self.elements_list.copy() + found = False + for i, element in enumerate(rotated_elements): + if element.name == name: + rotated_elements[i] = element.rotated(theta, rotate_beta) + found = True + + assert found, "Element to rotate was not found in the group" + + return ElementsGroup(rotated_elements, name=self.name, tag=self.tag, description=self.description) + + def rotated(self, theta: float, rotate_beta: bool = False) -> ElementsGroup: + """ + Returns a copy if the self-ElementsGroup which has been rotated counterclockwise in the transverse plane by an + angle of theta radians. If the rotate_beta flag is enable, the beta_x and beta_y values of the new ElementsGroup + are rotated correspondingly. + :param theta: The angle, in radians, the ElementsGroup is to be rotated by + :param rotate_beta: A flag indicating whether or not the beta_x and beta_y values of the new ElementsGroup + should also be rotated by theta. Enabled by default. + :return: A newly initialized copy of the self-ElementsGroup which has been rotated counterclockwise in the + transverse plane by an angle theta. + """ + rotated_elements = [] + + for element in rotated_elements: + rotated_elements.append(element.rotated(theta, rotate_beta)) + + return ElementsGroup(rotated_elements, name=self.name, tag=self.tag, description=self.description) + + def changed_betas(self, new_beta_x: float, new_beta_y: float) -> ElementsGroup: + elements_list_new = [] + for i_elem, elem in enumerate(self.elements_list): + elements_list_new.append(elem.changed_betas(new_beta_x, new_beta_y)) + + return ElementsGroup(elements_list_new, name=self.name, tag=self.tag, description=self.description) + + def get_element(self, name_string: str): + for element in self.elements_list: + if element.name == name_string: + return element + + raise KeyError(f"'{self.name}' has no element named '{name_string}'.") + diff --git a/xwakes/wit/interface.py b/xwakes/wit/interface.py new file mode 100644 index 00000000..623c9556 --- /dev/null +++ b/xwakes/wit/interface.py @@ -0,0 +1,813 @@ +import os + +from .component import Component +from .element import Element + +import subprocess +from typing import Tuple, List, Optional, Dict, Any, Union +from dataclasses import dataclass +from pathlib import Path +from hashlib import sha256 + +import numpy as np +from yaml import load, BaseLoader +from scipy.interpolate import interp1d + +# A dictionary mapping the datafile-prefixes (as used in IW2D) to (is_impedance, plane, (a, b, c, d)) +# Where is impedance is True if the component in question is an impedance component, and False if it is a +# wake component, and a, b, c and d are the source and test exponents of the component +component_names = {'wlong': (False, 'z', (0, 0, 0, 0)), + 'wxdip': (False, 'x', (1, 0, 0, 0)), + 'wydip': (False, 'y', (0, 1, 0, 0)), + 'wxqua': (False, 'x', (0, 0, 1, 0)), + 'wyqua': (False, 'y', (0, 0, 0, 1)), + 'wxcst': (False, 'x', (0, 0, 0, 0)), + 'wycst': (False, 'y', (0, 0, 0, 0)), + 'zlong': (True, 'z', (0, 0, 0, 0)), + 'zxdip': (True, 'x', (1, 0, 0, 0)), + 'zydip': (True, 'y', (0, 1, 0, 0)), + 'zxqua': (True, 'x', (0, 0, 1, 0)), + 'zyqua': (True, 'y', (0, 0, 0, 1)), + 'zxcst': (True, 'x', (0, 0, 0, 0)), + 'zycst': (True, 'y', (0, 0, 0, 0))} + +# The parent directory of this file +IW2D_SETTINGS_PATH = Path.home().joinpath('pywit').joinpath('config').joinpath('iw2d_settings.yaml') + + +def get_component_name(is_impedance, plane, exponents): + """ + Get the component name from is_impedance, plane and exponents (doing the + reverse operation of the dictionary in component_names) + :param is_impedance: True for impedance component, False for wake + :param plane: plane ('x', 'y' or 'z') + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :return: str with component name (e.g. 'zydip' or 'wxqua') + """ + comp_list = [comp_name for comp_name, v in component_names.items() + if v == (is_impedance, plane, exponents)] + if len(comp_list) != 1: + raise ValueError(f"({is_impedance},{plane},{exponents}) cannot be found in" + " the values of component_names dictionary") + + return comp_list[0] + + +def get_iw2d_config_value(key: str) -> Any: + with open(IW2D_SETTINGS_PATH, 'r') as file: + config = load(file, Loader=BaseLoader) + + return config[key] + + +def import_data_iw2d(directory: Union[str, Path], + common_string: str) -> List[Tuple[bool, str, Tuple[int, int, int, int], np.ndarray]]: + """ + Imports data on the format generated by the IW2D library and prepares it for construction of Components and + Elements in PyWIT + :param directory: The directory where the .dat files are located. All .dat files must be in the root of this + directory + :param common_string: A string preceding ".dat" in the filenames of all files to be imported + :return: A list of tuples, one for each imported file, on the form (is_impedance, plane, (a, b, c, d), data), + where data is a numpy array with 2 or 3 columns, one for each column of the imported datafile. + """ + # The tuples are iteratively appended to this array + component_recipes = [] + + # Keeps track of what combinations of (is_impedance, plane, exponents) have been imported to avoid duplicates + seen_configs = [] + + # A list of all of the filenames in the user-specified directory + filenames = os.listdir(directory) + for i, filename in enumerate(filenames): + # If the string preceding ".dat" in the filename does not match common_string, or if the first 5 letters + # of the filename are not recognized as a type of impedance/wake, the file is skipped + if filename[-4 - len(common_string):-4] != common_string or filename[:5].lower() not in component_names: + continue + + # The values of is_impedance, plane and exponents are deduced from the first 5 letters of the filename using + # the component_names-dictionary + is_impedance, plane, exponents = component_names[filename[:5].lower()] + + # Validates that the combination of (is_impedance, plane, exponents) is unique + assert (is_impedance, plane, exponents) not in seen_configs, \ + f"The {'impedance' if is_impedance else 'wake'} files " \ + f"'{filename}' and '{filenames[seen_configs.index((is_impedance, plane, exponents))]}' " \ + f"both correspond to the {plane}-plane with exponents {exponents}." + seen_configs.append((is_impedance, plane, exponents)) + + # Loads the data from the file as a numpy array + data = np.loadtxt(f"{directory}/{filename}", delimiter=" ", skiprows=1) + + # Appends the constructed tuple to component_recipes + component_recipes.append((is_impedance, plane, exponents, data)) + + # Validates that at least one file in the directory matched the user-specified common_string + assert component_recipes, f"No files in '{directory}' matched the common string '{common_string}'." + return component_recipes + + +def create_component_from_data(is_impedance: bool, plane: str, exponents: Tuple[int, int, int, int], + data: np.ndarray, relativistic_gamma: float, + extrapolate_to_zero: bool = False) -> Component: + """ + Creates a Component from a component recipe, e.g. as generated by import_data_iw2d + :param is_impedance: a bool which is True if the component to be generated is an impedance component, and False + if it is a wake component + :param plane: the plane of the component + :param exponents: the exponents of the component on the form (a, b, c, d) + :param data: a numpy-array with 2 or 3 columns corresponding to (frequency, Re[impedance], Im[impedance]) or + (position, Re[wake], Im[wake]), where the imaginary column is optional + :param relativistic_gamma: The relativistic gamma used in the computation of the data files. Necessary for + converting the position-data of IW2D into time-data for PyWIT + :param extrapolate_to_zero: a flag specifying if the data should be extrapolated to zero. If it is False (default + value), the data are extrapolated using the first and last value of the data. + :return: A Component object as specified by the input + """ + # Extracts the position/frequency column of the data array + x = data[:, 0] + + if not is_impedance: + # Converts position-data to time-data using Lorentz factor + x /= 299792458 * np.sqrt(1 - (1 / relativistic_gamma ** 2)) + + # Extracts the wake/values from the data array + y = data[:, 1] + (1j * data[:, 2] if data.shape[1] == 3 else 0) + + # Creates a callable impedance/wake function from the data array + if extrapolate_to_zero: + func = interp1d(x=x, y=y, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0, 0)) + else: + func = interp1d(x=x, y=y, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(y[0], y[-1])) + + # Initializes and returns a component based on the parameters provided + return Component(impedance=(func if is_impedance else None), + wake=(None if is_impedance else func), + plane=plane, + source_exponents=exponents[:2], + test_exponents=exponents[2:], ) + + +@dataclass(frozen=True, eq=True) +class Layer: + # The distance in mm of the inner surface of the layer from the reference orbit + thickness: float + dc_resistivity: float + resistivity_relaxation_time: float + re_dielectric_constant: float + magnetic_susceptibility: float + permeability_relaxation_frequency: float + + +@dataclass(frozen=True, eq=True) +class Sampling: + start: float + stop: float + # 0 = logarithmic, 1 = linear, 2 = both + scan_type: int + added: Tuple[float] + sampling_exponent: Optional[float] = None + points_per_decade: Optional[float] = None + min_refine: Optional[float] = None + max_refine: Optional[float] = None + n_refine: Optional[float] = None + + +# Define several dataclasses for IW2D input elements. We must split mandatory +# and optional arguments into private dataclasses to respect the resolution +# order. The public classes RoundIW2DInput and FlatIW2D input inherit from +# from the private classes. +# https://stackoverflow.com/questions/51575931/class-inheritance-in-python-3-7-dataclasses + +@dataclass(frozen=True, eq=True) +class _IW2DInputBase: + machine: str + length: float + relativistic_gamma: float + calculate_wake: bool + f_params: Sampling + + +@dataclass(frozen=True, eq=True) +class _IW2DInputOptional: + z_params: Optional[Sampling] = None + long_factor: Optional[float] = None + wake_tol: Optional[float] = None + freq_lin_bisect: Optional[float] = None + comment: Optional[str] = None + + +@dataclass(frozen=True, eq=True) +class IW2DInput(_IW2DInputOptional, _IW2DInputBase): + pass + + +@dataclass(frozen=True, eq=True) +class _RoundIW2DInputBase(_IW2DInputBase): + layers: Tuple[Layer] + inner_layer_radius: float + # (long, xdip, ydip, xquad, yquad) + yokoya_factors: Tuple[float, float, float, float, float] + + +@dataclass(frozen=True, eq=True) +class _RoundIW2DInputOptional(_IW2DInputOptional): + pass + + +@dataclass(frozen=True, eq=True) +class RoundIW2DInput(_RoundIW2DInputOptional, _RoundIW2DInputBase): + pass + + +@dataclass(frozen=True, eq=True) +class _FlatIW2DInputBase(_IW2DInputBase): + top_bottom_symmetry: bool + top_layers: Tuple[Layer] + top_half_gap: float + + +@dataclass(frozen=True, eq=True) +class _FlatIW2DInputOptional(_IW2DInputOptional): + bottom_layers: Optional[Tuple[Layer]] = None + bottom_half_gap: Optional[float] = None + + +@dataclass(frozen=True, eq=True) +class FlatIW2DInput(_FlatIW2DInputOptional, _FlatIW2DInputBase): + pass + + +def _iw2d_format_layer(layer: Layer, n: int) -> str: + """ + Formats the information describing a single layer into a string in accordance with IW2D standards. + Intended only as a helper-function for create_iw2d_input_file. + :param layer: A Layer object + :param n: The 1-indexed index of the layer + :return: A string on the correct format for IW2D + """ + return (f"Layer {n} DC resistivity (Ohm.m):\t{layer.dc_resistivity}\n" + f"Layer {n} relaxation time for resistivity (ps):\t{layer.resistivity_relaxation_time * 1e12}\n" + f"Layer {n} real part of dielectric constant:\t{layer.re_dielectric_constant}\n" + f"Layer {n} magnetic susceptibility:\t{layer.magnetic_susceptibility}\n" + f"Layer {n} relaxation frequency of permeability (MHz):\t{layer.permeability_relaxation_frequency / 1e6}\n" + f"Layer {n} thickness in mm:\t{layer.thickness * 1e3}\n") + + +def _iw2d_format_freq_params(params: Sampling) -> str: + """ + Formats the frequency-parameters of an IW2DInput object to a string in accordance with IW2D standards. + Intended only as a helper-function for create_iw2d_input_file. + :param params: Parameters specifying a frequency-sampling + :return: A string on the correct format for IW2D + """ + lines = [f"start frequency exponent (10^) in Hz:\t{np.log10(params.start)}", + f"stop frequency exponent (10^) in Hz:\t{np.log10(params.stop)}", + f"linear (1) or logarithmic (0) or both (2) frequency scan:\t{params.scan_type}"] + + if params.sampling_exponent is not None: + lines.append(f"sampling frequency exponent (10^) in Hz (for linear):\t{np.log10(params.sampling_exponent)}") + + if params.points_per_decade is not None: + lines.append(f"Number of points per decade (for log):\t{params.points_per_decade}") + + if params.min_refine is not None: + lines.append(f"when both, fmin of the refinement (in THz):\t{params.min_refine / 1e12}") + + if params.max_refine is not None: + lines.append(f"when both, fmax of the refinement (in THz):\t{params.max_refine / 1e12}") + + if params.n_refine is not None: + lines.append(f"when both, number of points in the refinement:\t{params.n_refine}") + + lines.append(f"added frequencies [Hz]:\t{' '.join(str(f) for f in params.added)}") + + return "\n".join(lines) + "\n" + + +def _iw2d_format_z_params(params: Sampling) -> str: + """ + Formats the position-parameters of an IW2DInput object to a string in accordance with IW2D standards. + Intended only as a helper-function for create_iw2d_input_file. + :param params: Parameters specifying a position-sampling + :return: A string on the correct format for IW2D + """ + lines = [f"linear (1) or logarithmic (0) or both (2) scan in z for the wake:\t{params.scan_type}"] + + if params.sampling_exponent is not None: + lines.append(f"sampling distance in m for the linear sampling:\t{params.sampling_exponent}") + + if params.min_refine is not None: + lines.append(f"zmin in m of the linear sampling:\t{params.min_refine}") + + if params.max_refine is not None: + lines.append(f"zmax in m of the linear sampling:\t{params.max_refine}") + + if params.points_per_decade is not None: + lines.append(f"Number of points per decade for the logarithmic sampling:\t{params.points_per_decade}") + + lines.append(f"exponent (10^) of zmin (in m) of the logarithmic sampling:\t{np.log10(params.start)}") + lines.append(f"exponent (10^) of zmax (in m) of the logarithmic sampling:\t{np.log10(params.stop)}") + lines.append(f"added z [m]:\t{' '.join(str(z) for z in params.added)}") + + return "\n".join(lines) + "\n" + + +def create_iw2d_input_file(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], filename: Union[str, Path]) -> None: + """ + Writes an IW2DInput object to the specified filename using the appropriate format for interfacing with the IW2D + software. + :param iw2d_input: An IW2DInput object to be written + :param filename: The filename (including path) of the file the IW2DInput object will be written to + :return: Nothing + """ + # Creates the input-file at the location specified by filename + file = open(filename, 'w') + + file.write(f"Machine:\t{iw2d_input.machine}\n" + f"Relativistic Gamma:\t{iw2d_input.relativistic_gamma}\n" + f"Impedance Length in m:\t{iw2d_input.length}\n") + + # Just pre-defining layers to avoid potentially unbound variable later on + layers = [] + if isinstance(iw2d_input, RoundIW2DInput): + file.write(f"Number of layers:\t{len(iw2d_input.layers)}\n" + f"Layer 1 inner radius in mm:\t{iw2d_input.inner_layer_radius * 1e3}\n") + layers = iw2d_input.layers + elif isinstance(iw2d_input, FlatIW2DInput): + if iw2d_input.bottom_layers: + print("WARNING: bottom layers of IW2D input object are being ignored because the top_bottom_symmetry flag " + "is enabled") + file.write(f"Number of upper layers in the chamber wall:\t{len(iw2d_input.top_layers)}\n") + if iw2d_input.top_layers: + file.write(f"Layer 1 inner half gap in mm:\t{iw2d_input.top_half_gap * 1e3}\n") + layers = iw2d_input.top_layers + + for i, layer in enumerate(layers): + file.write(_iw2d_format_layer(layer, i + 1)) + + if isinstance(iw2d_input, FlatIW2DInput) and not iw2d_input.top_bottom_symmetry: + file.write(f"Number of lower layers in the chamber wall:\t{len(iw2d_input.bottom_layers)}\n") + if iw2d_input.bottom_layers: + file.write(f"Layer -1 inner half gap in mm:\t{iw2d_input.bottom_half_gap * 1e3}\n") + for i, layer in enumerate(iw2d_input.bottom_layers): + file.write(_iw2d_format_layer(layer, -(i + 1))) + + if isinstance(iw2d_input, FlatIW2DInput): + file.write(f"Top bottom symmetry (yes or no):\t{'yes' if iw2d_input.top_bottom_symmetry else 'no'}\n") + + file.write(_iw2d_format_freq_params(iw2d_input.f_params)) + if iw2d_input.z_params is not None: + file.write(_iw2d_format_z_params(iw2d_input.z_params)) + + if isinstance(iw2d_input, RoundIW2DInput): + file.write(f"Yokoya factors long, xdip, ydip, xquad, yquad:\t" + f"{' '.join(str(n) for n in iw2d_input.yokoya_factors)}\n") + + for desc, val in zip(["factor weighting the longitudinal impedance error", + "tolerance (in wake units) to achieve", + "frequency above which the mesh bisecting is linear [Hz]", + "Comments for the output files names"], + [iw2d_input.long_factor, iw2d_input.wake_tol, iw2d_input.freq_lin_bisect, iw2d_input.comment]): + if val is not None: + file.write(f"{desc}:\t{val}\n") + + file.close() + + +def check_already_computed(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], + name: str) -> Tuple[bool, str, Union[str, Path]]: + """ + Checks if a simulation with inputs iw2d_input is already present in the hash database. + :param iw2d_input: an iw2d input object + :param name: the name of the object + :return already_computed: boolean indicating if the iw2d_inputs have been already + computed + :return input_hash: string with the hash key corresponding to the inputs + :return working_directory: the path to the directory where the files were put, which is built as + `// + """ + projects_path = Path(get_iw2d_config_value('project_directory')) + + # initialize read ready to all False for convenience + # create the hash key + input_hash = sha256(iw2d_input.__str__().encode()).hexdigest() + + # we have three levels of directories: the first two are given by the first and second letters of the hash keys, + # the third is given by the rest of the hash keys. + directory_level_1 = projects_path.joinpath(input_hash[0:2]) + directory_level_2 = directory_level_1.joinpath(input_hash[2:4]) + working_directory = directory_level_2.joinpath(input_hash[4:]) + + already_computed = True + + # check if the directories exist. If they do not exist we create + if not os.path.exists(directory_level_1): + already_computed = False + os.mkdir(directory_level_1) + + if not os.path.exists(directory_level_2): + already_computed = False + os.mkdir(directory_level_2) + + if not os.path.exists(working_directory): + already_computed = False + os.mkdir(working_directory) + + components = [] + if not iw2d_input.calculate_wake: + for component in component_names.keys(): + # the ycst component is only given in the case of a flat chamber and the x component is never given + if component.startswith('z') and 'cst' not in component: + components.append(component) + if isinstance(iw2d_input, FlatIW2DInput): + components.append('zycst') + else: + for component in component_names.keys(): + # if the wake is computed, all keys from component_names dict are added, except the constant impedance/wake + # in first instance. If the simulation is a flat chamber, we add the vertical constant impedance/wake + if 'cst' not in component: + components.append(component) + if isinstance(iw2d_input, FlatIW2DInput): + components.append('wycst') + components.append('zycst') + + # The simulation seems to have been already computed, but we check if all the components of the impedance + # wake have been computed. If not, the computation will be relaunched + if already_computed: + # this list also includes the input file but it doesn't matter + computed_components = [name[0:5].lower() for name in os.listdir(working_directory)] + + for component in components: + if component not in computed_components: + already_computed = False + break + + return already_computed, input_hash, working_directory + + +def check_valid_hash_chunk(hash_chunk: str, length: int): + """ + Checks that the hash_chunk string can be the part of an hash key of given length. This means that hash_chunk must + have the right length and it must be a hexadecimal string + :param hash_chunk: the string to be checked + :param length: the length which the hash chunk should have + :return: True if hash_chunk is valid, False otherwise + """ + if len(hash_chunk) != length: + return False + + # check if the hash is an hexadecimal string + try: + int(hash_chunk, 16) + return True + except ValueError: + return False + + +def check_valid_working_directory(working_directory: Path): + """ + Checks if working_directory is valid. To be valid working directory must be of the form + `/hash[0:2]/hash[2:4]/hash[4:]` + :param working_directory: the path to the directory to be checked + :return: True if the working_directory is valid, False otherwise + """ + projects_path = Path(get_iw2d_config_value('project_directory')) + + if working_directory.parent.parent.parent != projects_path: + raise ValueError(f"The working directory must be located inside {projects_path}") + + return (check_valid_hash_chunk(working_directory.parent.parent.name, 2) and + check_valid_hash_chunk(working_directory.parent.name, 2) and + check_valid_hash_chunk(working_directory.name, 60)) + + +def add_iw2d_input_to_database(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], input_hash: str, + working_directory: Union[str, Path]): + """ + Add the iw2d input to the repository containing the simulations + :param iw2d_input: the input object of the IW2D simulation + :param input_hash: the hash key corresponding to the input + :param working_directory: the directory where to put the iw2d input file + """ + if type(working_directory) == str: + working_directory = Path(working_directory) + + if not check_valid_working_directory(working_directory): + raise ValueError("working directory is not in the right format. The right format is " + "`/hash[0:2]/hash[2:4]/hash[4:]`") + + directory_level_1 = working_directory.parent.parent + directory_level_2 = working_directory.parent + + if not os.path.exists(directory_level_1): + os.mkdir(directory_level_1) + if not os.path.exists(directory_level_2): + os.mkdir(directory_level_2) + + working_directory = directory_level_2.joinpath(input_hash[4:]) + + if not os.path.exists(working_directory): + os.mkdir(working_directory) + + create_iw2d_input_file(iw2d_input, working_directory.joinpath(f"input.txt")) + + +def create_element_using_iw2d(iw2d_input: Union[FlatIW2DInput, RoundIW2DInput], name: str, beta_x: float, beta_y: float, + tag: str = 'IW2D', extrapolate_to_zero: bool = False) -> Element: + """ + Create and return an Element using IW2D object. + :param iw2d_input: the IW2DInput object + :param name: the name of the Element + :param beta_x: the beta function value in the x-plane at the position of the Element + :param beta_y: the beta function value in the x-plane at the position of the Element + :param tag: a tag string for the Element + :param extrapolate_to_zero: a flag specifying if the data should be extrapolated to zero. If it is False (default + value), the data are extrapolated using the first and last value of the data. + :return: The newly computed Element + """ + assert " " not in name, "Spaces are not allowed in element name" + + assert verify_iw2d_config_file(), "The binary and/or project directories specified in config/iw2d_settings.yaml " \ + "do not exist or do not contain the required files and directories." + + # when looking for this IW2DInput in the database, the comment and the machine name don't necessarily need to be + # the same as the in the old simulation so we ignore it for creating the hash + iw2d_input_dict = iw2d_input.__dict__ + comment = iw2d_input_dict['comment'] + machine = iw2d_input_dict['machine'] + iw2d_input_dict['comment'] = '' + iw2d_input_dict['machine'] = '' + + # the path to the folder containing the IW2D executables + bin_path = Path(get_iw2d_config_value('binary_directory')) + # the path to the folder containing the database of already computed elements + + # check if the element is already present in the database and create the hash key corresponding to the IW2D input + already_computed, input_hash, working_directory = check_already_computed(iw2d_input, name) + + if already_computed: + print(f"The computation of '{name}' has already been performed with the exact given parameters. " + f"These results will be used to generate the element.") + + # if an element with the same inputs is not found inside the database, perform the computations and add the results + # to the database + if not already_computed: + add_iw2d_input_to_database(iw2d_input, input_hash, working_directory) + bin_string = ("wake_" if iw2d_input.calculate_wake else "") + \ + ("round" if isinstance(iw2d_input, RoundIW2DInput) else "flat") + "chamber.x" + subprocess.run(f'{bin_path.joinpath(bin_string)} < input.txt', shell=True, cwd=working_directory) + + # When the wake is computed with IW2D, a second set of files is provided by IW2D. These correspond to a "converged" + # simulation with double the number of mesh points for the wake. They files have the _precise suffix to their name. + # If the wake is computed, we retrieve these file to create the wit element. + common_string = "_precise" if iw2d_input.calculate_wake else '' + + component_recipes = import_data_iw2d(directory=working_directory, common_string=common_string) + + iw2d_input_dict['comment'] = comment + iw2d_input_dict['machine'] = machine + + return Element(length=iw2d_input.length, + beta_x=beta_x, beta_y=beta_y, + components=[create_component_from_data(*recipe, relativistic_gamma=iw2d_input.relativistic_gamma, + extrapolate_to_zero=extrapolate_to_zero) + for recipe in component_recipes], + name=name, tag=tag, description='A resistive wall element created using IW2D') + + +def verify_iw2d_config_file() -> bool: + bin_path = Path(get_iw2d_config_value('binary_directory')) + projects_path = Path(get_iw2d_config_value('project_directory')) + if not bin_path.exists() or not projects_path.exists(): + return False + + contents = os.listdir(bin_path) + for filename in ('flatchamber.x', 'roundchamber.x', 'wake_flatchamber.x', 'wake_roundchamber.x'): + if filename not in contents: + return False + + return True + + +def _typecast_sampling_dict(d: Dict[str, str]) -> Dict[str, Any]: + added = [float(f) for f in d['added'].split()] if 'added' in d else [] + added = tuple(added) + scan_type = int(d['scan_type']) + d.pop('added'), d.pop('scan_type') + + new_dict = {k: float(v) for k, v in d.items()} + new_dict['added'] = added + new_dict['scan_type'] = scan_type + return new_dict + + +def _create_iw2d_input_from_dict(d: Dict[str, Any]) -> Union[FlatIW2DInput, RoundIW2DInput]: + is_round = d['is_round'].lower() in ['true', 'yes', 'y', '1'] + d.pop('is_round') + layers, inner_layer_radius, yokoya_factors = list(), float(), tuple() + top_layers, top_half_gap, bottom_layers, bottom_half_gap = list(), float(), None, None + + if is_round: + inner_layer_radius = d['inner_layer_radius'] + if 'layers' in d: + layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['layers']] + layers = [Layer(**kwargs) for kwargs in layers_dicts] + d.pop('layers') + else: + if 'top_layers' in d: + top_layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['top_layers']] + top_layers = [Layer(**kwargs) for kwargs in top_layers_dicts] + top_half_gap = d['top_half_gap'] + d.pop('top_layers') + if d['top_bottom_symmetry'].lower() in ['true', 'yes', 'y', '1']: + bottom_layers = None + else: + bottom_layers_dicts = [{k: float(v) for k, v in layer.items()} for layer in d['bottom_layers']] + bottom_layers = [Layer(**kwargs) for kwargs in bottom_layers_dicts] + bottom_half_gap = d['bottom_half_gap'] + d.pop('bottom_layers') + + if 'yokoya_factors' in d: + yokoya_factors = tuple(float(x) for x in d['yokoya_factors'].split()) + d.pop('yokoya_factors') + + f_params = Sampling(**_typecast_sampling_dict(d['f_params'])) + z_params = Sampling(**_typecast_sampling_dict(d['z_params'])) \ + if d['calculate_wake'].lower() in ['true', 'yes', 'y', '1'] else None + + d.pop('f_params') + d.pop('z_params', None) + + transformations = { + 'machine': str, + 'length': float, + 'relativistic_gamma': float, + 'calculate_wake': lambda x: x.lower() in ['true', 'yes', 'y', '1'], + 'long_factor': float, + 'wake_tol': float, + 'freq_lin_bisect': float, + 'comment': str + } + + new_dict = {k: transformations[k](d[k]) if k in d else None for k in transformations} + + if is_round: + return RoundIW2DInput( + f_params=f_params, + z_params=z_params, + layers=tuple(layers), + inner_layer_radius=inner_layer_radius, + yokoya_factors=yokoya_factors, + **new_dict + ) + else: + return FlatIW2DInput( + f_params=f_params, + z_params=z_params, + top_bottom_symmetry=d['top_bottom_symmetry'].lower() in ['true', 'yes', 'y', '1'], + top_layers=tuple(top_layers), + top_half_gap=top_half_gap, + bottom_layers=bottom_layers, + bottom_half_gap=bottom_half_gap, + **new_dict + ) + + +def create_iw2d_input_from_yaml(name: str) -> Union[FlatIW2DInput, RoundIW2DInput]: + """ + Create a IW2DInput object from one of the inputs specified in the `pywit/config/iw2d_inputs.yaml` database + :param name: the name of the input which is read from the yaml database + :return: the newly initialized IW2DInput object + """ + path = Path.home().joinpath('pywit').joinpath('config').joinpath('iw2d_inputs.yaml') + with open(path) as file: + inputs = load(file, Loader=BaseLoader) + d = inputs[name] + + return _create_iw2d_input_from_dict(d) + + +def create_multiple_elements_using_iw2d(iw2d_inputs: List[IW2DInput], names: List[str], + beta_xs: List[float], beta_ys: List[float]) -> List[Element]: + """ + Create and return a list of Element's using a list of IW2D objects. + :param iw2d_inputs: the list of IW2DInput objects + :param names: the list of names of the Element's + :param beta_xs: the list of beta function values in the x-plane at the position of each Element + :param beta_ys: the list of beta function values in the x-plane at the position of each Element + :return: the list of newly computed Element's + """ + assert len(iw2d_inputs) == len(names) == len(beta_xs) == len(beta_ys), "All input lists need to have the same" \ + "number of elements" + + for name in names: + assert " " not in name, "Spaces are not allowed in element name" + + assert verify_iw2d_config_file(), "The binary and/or project directories specified in config/iw2d_settings.yaml " \ + "do not exist or do not contain the required files and directories." + + + from joblib import Parallel, delayed + elements = Parallel(n_jobs=-1, prefer='threads')(delayed(create_element_using_iw2d)( + iw2d_inputs[i], + names[i], + beta_xs[i], + beta_ys[i] + ) for i in range(len(names))) + + return elements + + +def create_htcondor_input_file(iw2d_input: IW2DInput, name: str, directory: Union[str, Path]) -> None: + exec_string = "" + if iw2d_input.calculate_wake: + exec_string += "wake_" + exec_string += ("round" if isinstance(iw2d_input, RoundIW2DInput) else "flat") + "chamber.x" + + text = f"executable = {exec_string}\n" \ + f"input = {name}_input.txt\n" \ + f"ID = $(Cluster).$(Process)\n" \ + f"output = $(ID).out\n" \ + f"error = $(ID).err\n" \ + f"log = $(Cluster).log\n" \ + f"universe = vanilla\n" \ + f"initialdir = \n" \ + f"when_to_transfer_output = ON_EXIT\n" \ + f'+JobFlavour = "tomorrow"\n' \ + f'queue' + + with open(directory, 'w') as file: + file.write(text) + + +def _verify_iw2d_binary_directory(ignore_missing_files: bool = False) -> None: + bin_path = Path(get_iw2d_config_value('binary_directory')) + if not ignore_missing_files: + filenames = ('flatchamber.x', 'roundchamber.x', 'wake_flatchamber.x', 'wake_roundchamber.x') + assert all(filename in os.listdir(bin_path) for filename in filenames), \ + "In order to utilize IW2D with PyWIT, the four binary files 'flatchamber.x', 'roundchamber.x', " \ + f"'wake_flatchamber.x' and 'wake_roundchamber.x' (as generated by IW2D) must be placed in the directory " \ + f"'{bin_path}'." + + +def _read_cst_data(filename: Union[str, Path]) -> np.ndarray: + with open(filename, 'r') as f: + lines = f.readlines() + data = [] + for line in lines: + try: + data.append([float(e) for e in line.strip().split()]) + except ValueError: + pass + + return np.asarray(data) + + +def load_longitudinal_impedance_datafile(path: Union[str, Path]) -> Component: + data = _read_cst_data(path) + fs = data[:, 0] + zs = data[:, 1] + 1j * data[:, 2] + func = interp1d(x=fs, y=zs, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) + return Component(impedance=func, plane='z', source_exponents=(0, 0), test_exponents=(0, 0)) + + +def load_transverse_impedance_datafile(path: Union[str, Path]) -> Tuple[Component, Component, Component, Component]: + data = _read_cst_data(path) + fs = data[:, 0] + zs = [data[:, 2 * i + 1] + 1j * data[:, 2 * i + 2] for i in range(4)] + components = tuple() + for i, z in enumerate(zs): + exponents = [int(j == i) for j in range(4)] + func = interp1d(x=fs, y=z, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) + components += (Component(impedance=func, + plane='x' if i % 2 == 0 else 'y', + source_exponents=(exponents[0], exponents[1]), + test_exponents=(exponents[2], exponents[3])),) + + return components + + +def load_longitudinal_wake_datafile(path: Union[str, Path]) -> Component: + data = _read_cst_data(path) + ts = data[:, 0] + ws = data[:, 1] * 1e15 + func = interp1d(x=ts, y=ws, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) + return Component(wake=func, plane='z', source_exponents=(0, 0), test_exponents=(0, 0)) + + +def load_transverse_wake_datafile(path: Union[str, Path]) -> Tuple[Component, Component, Component, Component]: + data = _read_cst_data(path) + ts = data[:, 0] + ws = [data[:, i] * 1e15 for i in range(1, 5)] + components = tuple() + for i, w in enumerate(ws): + exponents = [int(j == i) for j in range(4)] + func = interp1d(x=ts, y=w, kind='linear', assume_sorted=True, bounds_error=False, fill_value=(0., 0.)) + components += (Component(wake=func, + plane='x' if i % 2 == 0 else 'y', + source_exponents=(exponents[0], exponents[1]), + test_exponents=(exponents[2], exponents[3])),) + + return components + diff --git a/xwakes/wit/landau_damping.py b/xwakes/wit/landau_damping.py new file mode 100644 index 00000000..97127573 --- /dev/null +++ b/xwakes/wit/landau_damping.py @@ -0,0 +1,155 @@ +import numpy as np +from scipy.special import exp1 +from scipy.optimize import newton + +from typing import Sequence, Tuple + + +def dispersion_integral_2d(tune_shift: np.ndarray, b_direct: float, b_cross: float, + distribution: str = 'gaussian'): + """ + Compute the dispersion integral in 2D from q complex tune shift, given the detuning coefficients (multiplied by + sigma). This is the integral vs Jx and Jy of Jx*dphi/dJx/(Q-bx*Jx-bxy*Jy-i0) (with phi the distribution function) + The transverse distribution can be 'gaussian' or 'parabolic'. + Note: for stability diagrams, use -1/dispersion_integral, and usually the convention is to plot -Im[Q] vs Re[Q]. + Reference: Berg-Ruggiero: https://cds.cern.ch/record/318826?ln=it + :param tune_shift: the complex tune shift + :param b_direct: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in + the x plane or $\alpha_y \sigma_y$ if working in the y plane) + :param b_cross: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in + the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) + :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' + :return: the dispersion integral + """ + if np.imag(tune_shift) == 0: + tune_shift = tune_shift - 1e-15 * 1j + + c = b_cross / b_direct + q = -tune_shift / b_direct + + if distribution == 'gaussian': + + i1 = (1 - c - (q + c - c * q) * np.exp(q) * exp1(q) + c * np.exp(q / c) * exp1(q / c)) / ((1 - c) ** 2) + + if np.isnan(i1): + i1 = 1. / q - (c + 2) / q ** 2 # asymptotic form for large q (assuming c is of order 1) + + elif distribution == 'parabolic': + + xi = q / 5. + + if np.abs(xi) > 100.: + # asymptotic form for large q (assuming c is of order 1) + i1 = 1. / q - (c + 2) / q ** 2 + else: + i1 = (((c + xi) ** 3 * np.log((1 + xi) / (c + xi)) + + (-1 + c) * (c * (c + 2 * c * xi + (-1 + 2 * c) * xi ** 2) + + (-1 + c) * xi ** 2 * (3 * c + xi + 2 * c * xi) * np.log(xi / (1 + xi)))) / + ((-1 + c) ** 2 * c ** 2)) + i1 = -i1 * 4. / 5. + # this is the same as in Scott Berg-Ruggiero CERN SL-AP-96-71 (AP) + + else: + raise ValueError("distribution can only be 'gaussian' or 'parabolic'") + + i = -i1 / b_direct + + # additional minus sign because we want the integral with dphi/dJx (derivative of distribution) on the + # numerator, so -[the one of Berg-Ruggiero] + return -i + + +def find_detuning_coeffs_threshold(tune_shift: complex, q_s: float, b_direct_ref: float, b_cross_ref: float, + fraction_of_qs_allowed_on_positive_side: float = 0.05, + distribution: str = 'gaussian', tolerance=1e-10): + """ + Compute the detuning coefficients (multiplied by sigma) corresponding to stability diagram threshold for a complex + tune shift. + It keeps fixed the ratio between b_direct_ref and b_cross_ref. + :param tune_shift: the tune shift for which the octupole threshold is computed + :param q_s: the synchrotron tune + :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in + the x plane or $\alpha_y \sigma_y$ if working in the y plane) + :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in + the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) + :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' + :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the + stability diagram), one can consider positive tune shift up to this fraction of q_s (default=5%) + :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding + and for the final check that the roots are actually proper roots. + :return: the detuning coefficients corresponding to the stability diagram threshold if the corresponding mode is + unstable, 0 if the corresponding mode is stable or np.nan if the threshold cannot be found (failure of Newton's + algorithm). + """ + # evaluate azimuthal mode number + l_mode = int(np.ceil(np.real(tune_shift) / q_s)) + if (l_mode - np.real(tune_shift) / q_s) > 1 - fraction_of_qs_allowed_on_positive_side: + l_mode -= 1 + # take away the shift from azimuthal mode number + tune_shift -= q_s * l_mode + + b_ratio = b_cross_ref/b_direct_ref + if tune_shift.imag < 0.: + + # function to solve (distance in imag. part w.r.t stab. diagram, as a function of oct. current) + def f(b_direct): + b_direct_i = b_direct + b_cross_i = b_ratio * b_direct + stab = [-1. / dispersion_integral_2d(t_s, b_direct_i, b_cross_i, distribution=distribution) for e in (-1, 1) + for t_s in b_direct_i * e * 10. ** np.arange(-3, 2, 0.01)[::e]] + # note: one has to reverse the table to get the interpolation right, for negative polarity (np.interp always + # wants monotonically increasing abscissae) + return tune_shift.imag - np.interp(tune_shift.real, np.real(stab)[::int(np.sign(b_direct_ref))], + np.imag(stab)[::int(np.sign(b_direct_ref))]) + + # Newton root finding + try: + b_direct_new = newton(f, b_direct_ref, tol=tolerance) + except RuntimeError: + b_direct_new = np.nan + else: + if np.abs(f(b_direct_new)) > tolerance: + b_direct_new = np.nan + else: + b_direct_new = 0. + + return b_direct_new, b_ratio*b_direct_new + + +def abs_first_item_or_nan(tup: Tuple): + if tup is not np.nan: + return abs(tup[0]) + else: + return np.nan + + +def find_detuning_coeffs_threshold_many_tune_shifts(tune_shifts: Sequence[complex], q_s: float, b_direct_ref: float, + b_cross_ref: float, distribution: str = 'gaussian', + fraction_of_qs_allowed_on_positive_side: float = 0.05, + tolerance=1e-10): + """ + Compute the detuning coefficients corresponding to the most stringent stability diagram threshold for a sequence of + complex tune shifts. It keeps fixed the ratio between b_direct_ref and b_cross_ref. + :param tune_shifts: the sequence of complex tune shifts + :param q_s: the synchrotron tune + :param b_direct_ref: the direct detuning coefficient multiplied by sigma (i.e. $\alpha_x \sigma_x$ if working in + the x plane or $\alpha_y \sigma_y$ if working in the y plane) + :param b_cross_ref: the cross detuning coefficient multiplied by sigma (i.e. $\alpha_{xy} \sigma_y$ if working in + the x plane or $\alpha_{yx} \sigma_x$ if working in the y plane) + :param distribution: the transverse distribution of the beam. It can be 'gaussian' or 'parabolic' + :param fraction_of_qs_allowed_on_positive_side: to determine azimuthal mode number l_mode (around which is drawn the + stability diagram), one can consider positive tuneshift up to this fraction of q_s (default=5%) + :param tolerance: tolerance on difference w.r.t stability diagram, for Newton's root finding + and for the final check that the roots are actually proper roots. + :return: the detuning coefficients corresponding to the most stringent stability diagram threshold for all the + given tune shifts if the corresponding mode is unstable, 0 if all modes are stable or np.nan if the + no threshold can be found (failure of Newton's algorithm). + """ + # find max octupole current required from a list of modes, given their tuneshifts + b_coefficients = np.array([find_detuning_coeffs_threshold( + tune_shift=tune_shift, q_s=q_s, b_direct_ref=b_direct_ref, + b_cross_ref=b_cross_ref, distribution=distribution, + fraction_of_qs_allowed_on_positive_side=fraction_of_qs_allowed_on_positive_side, + tolerance=tolerance) for tune_shift in tune_shifts if tune_shift is not np.nan]) + + return max(b_coefficients, key=abs_first_item_or_nan) diff --git a/pywit/materials.json b/xwakes/wit/materials.json similarity index 100% rename from pywit/materials.json rename to xwakes/wit/materials.json diff --git a/xwakes/wit/materials.py b/xwakes/wit/materials.py new file mode 100644 index 00000000..cdf73b2a --- /dev/null +++ b/xwakes/wit/materials.py @@ -0,0 +1,240 @@ +from .interface import Layer +from .utils import round_sigfigs + +from pathlib import Path +import numpy as np +import json +from typing import Callable, Tuple + + +def layer_from_dict(thickness: float, material_dict: dict) -> Layer: + """ + Define a layer from a dictionary containing the materials properties. + + :param thickness: layer thickness in m + :type thickness: float + :param material_dict: dictionary of material properties. 'dc_resistivity', 'resistivity_relaxation_time', + 're_dielectric_constant', 'magnetic_susceptibility', 'permeability_relaxation_frequency' are required + :type material_dict: dict + :return: Layer of the provided material + :rtype: Layer + """ + + # Check that the provided dict has the required entries to create a Layer object + # If not raise an AssertionError and indicate which property is missing + required_material_properties = np.array(['dc_resistivity', 'resistivity_relaxation_time', + 're_dielectric_constant', 'magnetic_susceptibility', + 'permeability_relaxation_frequency']) + # missing_properties_list is an array of bool. False indicates a missing key + missing_properties_list = np.array([key not in material_dict for key in required_material_properties]) + + assert not any(missing_properties_list), '{} missing from the input dictionary'.format( + ", ".join(required_material_properties[np.asarray(missing_properties_list)])) + + return Layer(thickness=thickness, + dc_resistivity=material_dict['dc_resistivity'], + resistivity_relaxation_time=material_dict['resistivity_relaxation_time'], + re_dielectric_constant=material_dict['re_dielectric_constant'], + magnetic_susceptibility=material_dict['magnetic_susceptibility'], + permeability_relaxation_frequency=material_dict['permeability_relaxation_frequency']) + + +def layer_from_json_material_library(thickness: float, material_key: str, + library_path: Path = Path(__file__).parent.joinpath('materials.json')) -> Layer: + """ + Define a layer using the materials.json library of materials properties. + + :param thickness: layer thickness in m + :type thickness: float + :param material_key: material key in the materials.json file + :type material_key: str + :param library_path: material library path, defaults to materials.json present in wit + :type library_path: Path, optional + :return: Layer of the selected material + :rtype: Layer + """ + + materials_library_json = json.loads(library_path.read_bytes()) + + # Check that the material is in the provided materials library + assert material_key in materials_library_json.keys(), f"Material {material_key} is not in library {library_path}" + + # Put the material properties in a dict + # The decoder will have converted NaN, Infinity and -Infinity from JSON to python nan, inf and -inf + # The entries in the materials library can contain additional fields (comment, reference) that are not + # needed for the Layer object creation. + material_properties_dict = materials_library_json[material_key] + + layer = layer_from_dict(thickness=thickness, material_dict=material_properties_dict) + + return layer + + +# Resistivity rho at B=0 vs. temperature and RRR +def rho_vs_T_Hust_Lankford(T: float, rho273K: float, RRR: float, + P: Tuple[float, float, float, float, float, float, float], + rhoc: Callable[[float], float] = lambda T: 0) -> float: + """ + Define a resistivity versus temperature and RRR law as found in Hust & Lankford, "Thermal + conductivity of aluminum, copper, iron, and tungsten for temperatures from 1K to the melting point", National Bureau + of Standards, 1984, Eqs. (1.2.3) to (1.2.6) p. 8. + Note that typically P4 given there has the wrong sign - for Cu and W at least one would otherwise + not get the right behaviour at high temperature - most probably there is a typo in the reference. + + :param T: temperature in K + :type T: float + :param rho273K: resistivity at 273 K in Ohm.m + :type rho273K: float + :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0 + :type RRR: float + :param P: tuple of the fitting coeeficient. The coefficient can be found in the reference above for copper (p. 22), + aluminum (p. 92), iron (p. 145) and tungsten (p. 204) + :type P: tuple + :param rhoc: function of temperature accounting for the residual deviation from the law in Ohm.m, defaults to 0. + This can be used with iron and tungsten for which the residual function is provided in the reference + :type rhoc: float + :return: the resistivity value, in Ohm.m + :rtype: float + """ + + # To follow the notation used in the book (coeficients from P1 to P7), we unpack the P tuple to new variables + P1 = P[0] + P2 = P[1] + P3 = P[2] + P4 = P[3] + P5 = P[4] + P6 = P[5] + P7 = P[6] + rho0 = rho273K/RRR + rhoi = P1*T**P2 / (1. + P1*P3*T**(P2+P4)*np.exp(-(P5/T)**P6)) + rhoc(T) + rhoi0 = P7*rhoi*rho0 / (rhoi+rho0) + + return (rho0+rhoi+rhoi0) + + +# magnetoresistance law (Kohler-like) drho/rho = f(B*rho_273K(B=0)/rho_T(B=0)) +def magnetoresistance_Kohler(B_times_Sratio: float, P: Tuple[float, ...]) -> float: + """ + Define a magnetoresistance law in the form drho/rho = f(B*rho_273K(B=0)/rho_T(B=0)) + + :param B_times_Sratio: product of magnetic field B and Sratio = rho_273K(B=0)/rho_T(B=0) at a given temperature T, in Tesla + :type B_times_Sratio: float + :param P: tuple of the Kohler curve fitting coefficients. Kohler curve are represented in log-log scales: if x and y are read on the + curve, the P coefficients come from the fitting of log10(x) and log10(y) + :type P: tuple of floats + :return: drho/rho the resistivity variation at a given magnetic field + :rtype: float + """ + + if B_times_Sratio == 0. or P == (0,): + return 0. + else: + return 10.**np.polyval(P, np.log10(B_times_Sratio)) + + +def copper_at_temperature(thickness: float, T: float = 300, RRR: float = 70, B: float = 0) -> Layer: + """ + Define a layer of pure copper material at any temperature, any B field and any RRR. + We use a magnetoresistance law fitted from the UPPER curve of the plot in NIST, "Properties of copper and copper + alloys at cryogenic temperatures", by Simon, Crexler and Reed, 1992 (p. 8-27, Fig. 8-14). + The upper curve was chosen, as it is in relative agreement with C. Rathjen measurements and actual LHC beam screens + (CERN EDMS document Nr. 329882). + The resistivity vs. temperature and RRR is found from Hust & Lankford (see above). + The law vs. temperature was found in good agreement with the NIST reference above (p. 8-5, Fig. 8-2). + + :param thickness: material thickness in m + :type thickness: float + :param T: temperature in K, defaults to 300 + :type T: float, optional + :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0, defaults to 70 + :type RRR: float, optional + :param B: magnetic field in T, defaults to 0 + :type B: float, optional + :return: a Layer object + :rtype: Layer + """ + + rho273K = 15.5*1e-9 # resistivity at 273K, in Ohm.m + + # resistivity vs temperature law coefficients, found in p. 22 of Hust and Lankford + # Here P4 = -1.14 (instead of +1.14 in the book) to get the proper law behavior + P = (1.171e-17, 4.49, 3.841e10, -1.14, 50., 6.428, 0.4531) + + # Coefficients for the magnetoresistance law, from Simon, Crexler and Reed (p. 8-27) + kohler_P = (0.029497104404715, 0.905633738689341, -2.361415783729567) + + rhoDC_B0 = rho_vs_T_Hust_Lankford(T, rho273K, RRR, P) # resistivity for B=0 + Sratio = rho_vs_T_Hust_Lankford(273, rho273K, RRR, P) / rhoDC_B0 + dc_resistivity = round_sigfigs(rhoDC_B0 * (1.+magnetoresistance_Kohler(B*Sratio, kohler_P)),3) # we round it to 3 significant digits + + # tauAC formula from Ascroft-Mermin (Z=1 for copper), Drude model (used also for other + # materials defined above) with parameters from CRC - Handbook of Chem. and Phys. + me = 9.10938e-31 # electron mass + e = 1.60218e-19 # electron elementary charge + rho_m = 9.0 # Cu volumic mass in g/cm3 (it is actually 9.02 at 4K, 8.93 at 273K, + # see https://www.copper.org/resources/properties/cryogenic/ ) + A = 63.546 # Cu atomic mass in g/mol + Z = 1 # number of valence electrons + n = 6.022e23*Z*rho_m*1e6 / A + tauAC = round_sigfigs(me / (n*dc_resistivity*e**2),3) # relaxation time (s) (3 significant digits) + + return Layer(thickness=thickness, + dc_resistivity=dc_resistivity, + resistivity_relaxation_time=tauAC, + re_dielectric_constant=1, + magnetic_susceptibility=0, + permeability_relaxation_frequency=np.inf) + + +def tungsten_at_temperature(thickness: float, T: float = 300, RRR: float = 70, B: float = 0) -> Layer: + """ + Define a layer of tungsten at any temperature, any B field and any RRR. + The resistivity vs. temperature and RRR is found from Hust & Lankford (see above). + The magnetoresistance effect is not included yet. + + :param thickness: material thickness in m + :type thickness: float + :param T: temperature in K, defaults to 300 + :type T: float, optional + :param RRR: residual resistivity ratio, i.e. rho(273K)/rho(0K) at B=0, defaults to 70 + :type RRR: float, optional + :param B: magnetic field in T, defaults to 0 + :type B: float, optional + :return: a Layer object + :rtype: Layer + """ + + rho273K = 48.4*1e-9 # resistivity at 273K, in Ohm.m + + # resistivity vs temperature law coefficients, found in p. 204 of Hust and Lankford + # Here P4 = -1.22 (instead of +1.22 in the book) to get the proper law behavior + P = (4.801e-16, 3.839, 1.88e10, -1.22, 55.63, 2.391, 0.0) + + # Residual equation for tungsten is found in p. 204 + # The correction is small and affects only the second or third decimal of the resistivity. + # We must therefore increase the number of decimal to see its effect + rhoc = lambda T: 0.7e-8 * np.log(T/560) * np.exp(-(np.log(T/1000)/0.6)**2) + + # Coefficients for the magnetoresistance law. Put to zero for now + kohler_P = (0,) + + rhoDC_B0 = rho_vs_T_Hust_Lankford(T, rho273K, RRR, P, rhoc) # resistivity for B=0 + Sratio = rho_vs_T_Hust_Lankford(273, rho273K, RRR, P, rhoc) / rhoDC_B0 + dc_resistivity = round_sigfigs(rhoDC_B0 * (1.+magnetoresistance_Kohler(B*Sratio, kohler_P)),3) # 3 significant digits + + # tauAC formula from Ascroft-Mermin (Z=2 for tungsten), Drude model (used also for other + # materials defined above) with parameters from CRC - Handbook of Chem. and Phys. + me = 9.10938e-31 # electron mass + e = 1.60218e-19 # electron elementary charge + rhom = 19.3 # W volumic mass in g/cm3 + A = 183.84 # W atomic mass in g/mol + Z = 2 # number of valence electrons + n = 6.022e23 * Z * rhom * 1e6 / A + tauAC = round_sigfigs(me/(n*dc_resistivity*e**2),3) # relaxation time (s) (3 significant digits) + + return Layer(thickness=thickness, + dc_resistivity=dc_resistivity, + resistivity_relaxation_time=tauAC, + re_dielectric_constant=1, + magnetic_susceptibility=0, + permeability_relaxation_frequency=np.inf) diff --git a/xwakes/wit/model.py b/xwakes/wit/model.py new file mode 100644 index 00000000..6e029e37 --- /dev/null +++ b/xwakes/wit/model.py @@ -0,0 +1,29 @@ +from .element import Element + +from typing import List, Optional, Tuple + + +class Model: + """ + Suggestion for structure of Model class + """ + def __init__(self, elements: List[Element] = None, lumped_betas: Optional[Tuple[float, float]] = None): + assert elements, "Model object needs to be initialized with at least one Element" + if lumped_betas is not None: + elements = [element.changed_betas(*lumped_betas) for element in elements] + + self.__elements = elements + self.__lumped_betas = lumped_betas + + @property + def elements(self): + return self.__elements + + @property + def total(self): + return sum(self.__elements) + + def append_element(self, element: Element): + if self.__lumped_betas is not None: + element = element.changed_betas(*self.__lumped_betas) + self.__elements.append(element) diff --git a/xwakes/wit/parameters.py b/xwakes/wit/parameters.py new file mode 100644 index 00000000..8dfe3b8f --- /dev/null +++ b/xwakes/wit/parameters.py @@ -0,0 +1,11 @@ +REL_TOL = 1e-8 +ABS_TOL = 1e-8 +PRECISION_FACTOR = 100. + +MIN_FREQ = 1e3 +MAX_FREQ = 1e10 +FREQ_P_FACTOR = 0.2 + +MIN_TIME = 1e-10 +MAX_TIME = 1e-7 +TIME_P_FACTOR = 0.2 diff --git a/xwakes/wit/plot.py b/xwakes/wit/plot.py new file mode 100644 index 00000000..ecdbbdcc --- /dev/null +++ b/xwakes/wit/plot.py @@ -0,0 +1,235 @@ +from .component import Component +from .element import Element +from .model import Model +from .parameters import * + +from typing import List, Dict, Union, Optional, Set +from collections import defaultdict + +import numpy as np + + +def plot_component(component: Component, plot_impedance: bool = True, plot_wake: bool = True, start: float = 1, + stop: float = 10000, points: int = 200, step_size: float = None, plot_real: bool = True, + plot_imag: bool = True) -> None: + """ + Function for plotting real and imaginary parts of impedance and wake functions of a single component + :param component: The component to be plotted + :param plot_impedance: A flag indicating if impedance should be plotted + :param plot_wake: A flag indicating if wake should be plotted + :param start: The first value on the x-axis of the plot + :param stop: The last value on the x-axis of the plot + :param points: The number of points to be evaluated for each line of the plot (alternative to step_size) + :param step_size: The distance between each point on the x-axis (alternative to points) + :param plot_real: A flag indicating if the real values should be plotted + :param plot_imag: A flag indicating if the imaginary values should be plotted + :return: Nothing + """ + import matplotlib.pyplot as plt + assert (plot_wake or plot_impedance) and (plot_real or plot_imag), "There is nothing to plot" + assert stop - start > 0, "stop must be greater than start" + if step_size: + assert step_size > 0, "Negative step_size not possible" + xs = np.arange(start, stop, step_size, dtype=float) + else: + xs = np.linspace(start, stop, points) + + legends = [] + if plot_impedance: + ys = component.impedance(xs) + if plot_real: + plt.plot(xs, ys.real) + legends.append("Re[Z(f)]") + if plot_imag: + plt.plot(xs, ys.imag) + legends.append("Im[Z(f)]") + if plot_wake: + ys = component.wake(xs) + if plot_real: + plt.plot(xs, ys.real) + legends.append("Re[W(z)]") + if plot_imag: + plt.plot(xs, ys.imag) + legends.append("Im[W(z)]") + + plt.legend(legends) + plt.show() + + +def plot_element_in_plane(element: Element, plane: str, plot_impedance: bool = True, plot_wake: bool = True, + start: float = 1, stop: float = 10000, points: int = 200, step_size: float = None, + plot_real: bool = True, plot_imag: bool = True): + """ + Sums all Components of Element in specified plane and plots the resulting Component with the parameters given. + :param element: The Element with Components to be plotted + :param plane: The plane in which the Components to be plotted lie + :param plot_impedance: A flag indicating if impedance should be plotted + :param plot_wake: A flag indicating if wake should be plotted + :param start: The first value on the x-axis of the plot + :param stop: The last value on the x-axis of the plot + :param points: The number of points to be evaluated for each line of the plot (alternative to step_size) + :param step_size: The distance between each point on the x-axis (alternative to points) + :param plot_real: A flag indicating if the real values should be plotted + :param plot_imag: A flag indicating if the imaginary values should be plotted + :return: Nothing + """ + plane = plane.lower() + assert plane in ['x', 'y', 'z'], f"{plane} is not a valid plane. Must be 'x', 'y' or 'z'" + component = sum([c for c in element.components if c.plane == plane]) + + assert component, f"Element has no components in plane {plane}" + plot_component(component, plot_impedance=plot_impedance, plot_wake=plot_wake, start=start, stop=stop, + points=points, step_size=step_size, plot_real=plot_real, plot_imag=plot_imag) + + +def plot_component_impedance(component: Component, logscale_x: bool = True, logscale_y: bool = True, + points: int = 1000, start=MIN_FREQ, stop=MAX_FREQ, title: Optional[str] = None) -> None: + import matplotlib.pyplot as plt + fig: plt.Figure = plt.figure() + ax: plt.Axes = fig.add_subplot(111) + fs = np.geomspace(start, stop, points) + impedances = component.impedance(fs) + reals, imags = impedances.real, impedances.imag + ax.plot(fs, reals, label='real') + ax.plot(fs, imags, label='imag') + ax.set_xlim(start, stop) + if title: + plt.title(title) + + if logscale_x: + ax.set_xscale('log') + + if logscale_y: + ax.set_yscale('log') + + ax.set_xlabel('Frequency [Hz]') + ax.set_ylabel('Z [Ohm / m]') + + plt.legend() + plt.show() + + +def plot_component_wake(component: Component, logscale_x: bool = True, logscale_y: bool = True, + points: int = 1000, start=MIN_TIME, stop=MAX_TIME, title: Optional[str] = None) -> None: + import matplotlib.pyplot as plt + fig: plt.Figure = plt.figure() + ax: plt.Axes = fig.add_subplot(111) + ts = np.geomspace(start, stop, points) + ax.plot(ts, component.wake(ts)) + ax.set_xlim(start, stop) + + ax.set_xlabel('Time [s]') + ax.set_ylabel('Wake [V / C * m]') + + if title: + plt.title(title) + + if logscale_x: + ax.set_xscale('log') + + if logscale_y: + ax.set_yscale('log') + + plt.show() + + +def generate_contribution_plots(model: Model, start_freq: float = MIN_FREQ, stop_freq: float = MAX_FREQ, + start_time: float = MIN_TIME, stop_time: float = MAX_TIME, points: int = 1000, + freq_scale: str = 'log', time_scale: str = 'log', absolute: bool = False) -> None: + + import matplotlib.pyplot as plt + + # TODO: use roi's to generate grid + fs = np.geomspace(start_freq, stop_freq, points) + ts = np.geomspace(start_time, stop_time, points) + + all_tags = set([e.tag for e in model.elements]) + elements: Dict[str, Union[int, Element]] = defaultdict(int) + + for element in model.elements: + elements[element.tag] += element + + components_defined_for_tag: Dict[str, Set[str]] = dict() + all_type_strings = set() + for tag, element in elements.items(): + components_defined_for_tag[tag] = {c.get_shorthand_type() for c in element.components} + all_type_strings.update(components_defined_for_tag[tag]) + + tags = list(all_tags) + cumulative_elements: List[Element] = [] + defined_type_strings: List[set] = [] + current_defined = set() + current_element = 0 + for tag in tags: + new_element = current_element + elements[tag] + current_defined = current_defined.union(components_defined_for_tag[tag]) + defined_type_strings.append(current_defined) + cumulative_elements.append(new_element) + current_element = new_element + + for type_string in all_type_strings: + wakes = np.asarray([np.zeros(shape=ts.shape)]) + real_impedances = np.asarray([np.zeros(shape=fs.shape)]) + imag_impedances = np.asarray([np.zeros(shape=fs.shape)]) + for i, element in enumerate(cumulative_elements): + if type_string not in defined_type_strings[i]: + wakes = np.vstack((wakes, wakes[-1])) + real_impedances = np.vstack((real_impedances, real_impedances[-1])) + imag_impedances = np.vstack((imag_impedances, imag_impedances[-1])) + continue + + component = [c for c in element.components if c.get_shorthand_type() == type_string][0] + if component.wake is not None: + array = component.wake(ts) + wakes = np.vstack((wakes, array)) + else: + wakes = np.vstack((wakes, wakes[-1])) + + if component.impedance is not None: + impedances = component.impedance(fs) + real_impedances = np.vstack((real_impedances, impedances.real)) + imag_impedances = np.vstack((imag_impedances, impedances.imag)) + else: + real_impedances = np.vstack((real_impedances, real_impedances[-1])) + imag_impedances = np.vstack((imag_impedances, imag_impedances[-1])) + + titles = (f'Re[Z] contribution - {type_string}', + f'Im[Z] contribution - {type_string}', + f'Wake contribution - {type_string}') + + for i, (array, xs, title) in enumerate(zip((real_impedances, imag_impedances, wakes), + (fs, fs, ts), titles)): + if sum(array[-1]) == 0: + continue + + if not absolute: + array = np.divide(array, array[-1]) + + fig: plt.Figure = plt.figure() + ax: plt.Axes = fig.add_subplot(111) + for j in range(len(all_tags) - 1, -1, -1): + ax.fill_between(xs, array[j], array[j + 1], label=tags[j]) + + ax.set_xscale(time_scale if i == 2 else freq_scale) + if not absolute: + ax.set_ylim(0, 1) + ax.set_xlim(start_time if i == 2 else start_freq, stop_time if i == 2 else stop_freq) + plt.title(title) + plt.legend() + plt.show() + + +def plot_total_impedance_and_wake(model: Model, logscale_y=True, start_freq: float = MIN_FREQ, + stop_freq: float = MAX_FREQ, start_time: float = MIN_TIME, + stop_time: float = MAX_TIME, points: int = 1000, + logscale_freq: bool = True, logscale_time: bool = True): + element = sum(model.elements) + for component in element.components: + if component.impedance: + plot_component_impedance(component, logscale_y=logscale_y, + title=f"Total impedance - {component.get_shorthand_type()}", start=start_freq, + stop=stop_freq, points=points, logscale_x=logscale_freq) + if component.wake: + plot_component_wake(component, logscale_y=logscale_y, + title=f"Total wake - {component.get_shorthand_type()}", start=start_time, + stop=stop_time, points=points, logscale_x=logscale_time) diff --git a/xwakes/wit/sacherer_formula.py b/xwakes/wit/sacherer_formula.py new file mode 100644 index 00000000..88e8eb99 --- /dev/null +++ b/xwakes/wit/sacherer_formula.py @@ -0,0 +1,202 @@ +import numpy as np +import sys + +from scipy.constants import e as q_p, m_p, c + +from typing import List, Callable, Iterable, Tuple, Union + + +def sacherer_formula(qp: float, nx_array: np.array, bunch_intensity: float, omegas: float, n_bunches: int, + omega_rev: float, tune: float, gamma: float, eta: float, bunch_length_seconds: float, m_max: int, + impedance_table: np.array = None, freq_impedance_table: np.array = None, + impedance_function: Callable[[float], float] = None, m0: float = m_p, charge: float = q_p, + mode_type: str = 'sinusoidal') -> Tuple[np.array, float, np.array]: + + r""" + Computes frequency shift and effective impedance from Sacherer formula, in transverse, in the case of low + intensity perturbations (no mode coupling), for modes of kind 'mode_type'. + + Documentation: see Elias Metral's USPAS 2009 course : Bunched beams transverse coherent + instabilities. + + NOTE: this is NOT the original Sacherer formula, which assumes an impedance normalized by beta + (see E. Metral, USPAS 2009 lectures, or C. Zannini, + https://indico.cern.ch/event/766028/contributions/3179810/attachments/1737652/2811046/Z_definition.pptx) + Here this formula is instead divided by beta (compared to Sacherer initial one), + so is valid with our usual definition of impedance (not beta-normalized). + + :param qp: the chromaticity (defined as $\frac{\Delta q \cdot p}{\Delta p}$ + :param nx_array: a vector of coupled bunch modes for which the tune shift is computed (it must contain integers in + the range (0, M-1)) + :param bunch_intensity: number of particles per bunch + :param omegas: the synchrotron angular frequency (i.e. $Q_s \cdot \omega_{rev}$) + :param n_bunches: the number of bunches + :param omega_rev: the revolution angular frequency (i.e. $2\cdot \pi f_{rev}$) + :param tune: machine tune in the considered plane (the TOTAL tune, including the integer part, must be passed) + :param gamma: the relativistic gamma factor of the beam + :param eta: the slippage factor (i.e. alpha_p - 1/gamma^2, with alpha_p the momentum compaction factor) + :param bunch_length_seconds: the total bunch length in seconds (4 times $\sigma$ for a Gaussian bunch) + :param m_max: specifies the range (-m_max to m_max) of the azimuthal modes to be considered + :param impedance_table: a numpy array giving the complex impedance at a discrete set of points. It must be specified + if impedance_function is not specified + :param freq_impedance_table: the frequencies at which the impedance is sampled. It must be specified if + impedance_function is not specified + :param impedance_function: the impedance function. It must be specified if impedance_table is not specified + :param m0: the rest mass of the considered particles + :param charge: the charge of the considered particles + :param mode_type: the type of modes in the effective impedance. It can be 'sinusoidal' (typically + well-adpated for protons) or 'hermite' (typically better for leptons). + + :return tune_shift_nx: tune shifts for all multibunch modes and synchrotron modes. It is an array of dimensions + ( len(nx_scan), (2*m_max+1) ) + :return tune_shift_m0: tune shift of the most unstable coupled-bunch mode with m=0 + :return effective_impedance: the effective impedance for all multibunch modes and synchrotron modes. It is an array + of dimensions ( len(nx_scan), (2*m_max+1) ) + """ + + def hmm(m_mode: int, omega: Union[float, np.ndarray]): + """ + Compute hmm power spectrum of Sacherer formula, for azimuthal mode number m, + at angular frequency 'omega' (rad/s) (can be an array), for total bunch length + 'bunch_length_seconds' (s), and for a kind of mode specified by 'mode_type' + (which can be 'hermite' or 'sinusoidal') + :param m_mode: the azimuthal mode number + :param omega: the angular frequency at which hmm is computed + """ + + if mode_type.lower().startswith('sinus'): + # best for protons + hmm_val = (((bunch_length_seconds * (np.abs(m_mode) + 1.)) ** 2 / (2. * np.pi ** 4)) * + (1. + (-1) ** m_mode * np.cos(omega * bunch_length_seconds)) / + (((omega * bunch_length_seconds / np.pi) ** 2 - (np.abs(m_mode) + 1.) ** 2) ** 2)) + + elif mode_type.lower() == 'hermite': + # best for leptons + hmm_val = (omega * bunch_length_seconds / 4) ** (2 * m_mode) * np.exp( + -(omega * bunch_length_seconds / 4.) ** 2) + + else: + raise ValueError("mode_type can only be 'sinusoidal' or 'hermite'") + + return hmm_val + + def hmm_weighted_sum(m_mode: int, nx_mode: int, weight_function: Callable[[float], complex] = None): + """ + Compute sum of hmm functions in the Sacherer formula, optionally + weighted by weight_function. + Note: In the end the sum runs over k with hmm taken at the angular frequencies + (k_offset+k*n_bunches)*omega0+m*omegas-omegaksi but the impedance is taken at + (k_offset+k*n_bunches)*omega0+m*omegas + :param m_mode: the azimuthal mode number + :param nx_mode: the coupled-bunch mode number + :param weight_function: function of frequency (NOT angular) giving + the sum weights (typically, it is the impedance) (optional) + :return: the (possibly weigthed) sum of hmm functions + """ + eps = 1.e-5 # relative precision of the summations + k_max = 20 + k_offset = nx_mode + fractional_tune + # sum initialization + omega_k = k_offset * omega_rev + m_mode * omegas + hmm_k = hmm(m_mode, omega_k - omega_ksi) + + omega = np.arange(-100.01 / bunch_length_seconds, 100.01 / bunch_length_seconds, + 0.01 / bunch_length_seconds) + + if weight_function is not None: + z_pk = weight_function(omega_k / (2 * np.pi)) + else: + z_pk = np.ones_like(omega_k) + + sum1_inner = z_pk * hmm_k + + k = np.arange(1, k_max + 1) + old_sum1 = 10. * sum1_inner + + while ((np.abs(np.real(sum1_inner - old_sum1))) > eps * np.abs(np.real(sum1_inner))) or ( + (np.abs(np.imag(sum1_inner - old_sum1))) > eps * np.abs(np.imag(sum1_inner))): + old_sum1 = sum1_inner + # omega_k^x and omega_-k^x in Elias's slides: + omega_k = (k_offset + k * n_bunches) * omega_rev + m_mode * omegas + omega_mk = (k_offset - k * n_bunches) * omega_rev + m_mode * omegas + # power spectrum function h(m,m) for k and -k: + hmm_k = hmm(m_mode, omega_k - omega_ksi) + hmm_mk = hmm(m_mode, omega_mk - omega_ksi) + + if weight_function is not None: + z_pk = weight_function(omega_k / (2 * np.pi)) + z_pmk = weight_function(omega_mk / (2 * np.pi)) + else: + z_pk = np.ones_like(omega_k) + z_pmk = np.ones_like(omega_mk) + + # sum + sum1_inner = sum1_inner + np.sum(z_pk * hmm_k) + np.sum(z_pmk * hmm_mk) + + k = k + k_max + + sum1_inner = np.squeeze(sum1_inner) # return a scalar if only one element + + return sum1_inner + + if impedance_function is not None and impedance_table is not None: + raise ValueError('Only one between impedance_function and impedance_table can be specified') + + if impedance_table is not None and freq_impedance_table is None: + raise ValueError('When impedance_table is specified, also the corresponding frequencies must be specified in' + 'omega_impedance_table') + + # some parameters + beta = np.sqrt(1. - 1. / (gamma ** 2)) # relativistic velocity factor + f0 = omega_rev / (2. * np.pi) # revolution angular frequency + single_bunch_current = charge * bunch_intensity * f0 # single-bunch current + fractional_tune = tune - np.floor(tune) # fractional part of the tune + bunch_length_seconds_meters = bunch_length_seconds * beta * c # full bunch length (in meters) + + if impedance_table is not None: + def impedance_function(x): + if np.isscalar(x): + x = np.array([x]) + ind_p = x >= 0 + ind_n = x < 0 + result = np.zeros_like(x, dtype=complex) + result[ind_p] = np.interp(x[ind_p], freq_impedance_table, impedance_table) + result[ind_n] = -np.interp(np.abs(x[ind_n]), freq_impedance_table, impedance_table).conjugate() + + return result + + tune_shift_nx = np.zeros((len(nx_array), 2 * m_max + 1), dtype=complex) + tune_shift_m0 = complex(0) + effective_impedance = np.zeros((len(nx_array), 2 * m_max + 1), dtype=complex) + + omega_ksi = qp * omega_rev / eta + + for inx, nx in enumerate(nx_array): # coupled-bunch modes + + for im, m in enumerate(range(-m_max, m_max + 1)): + # consider each synchrotron mode individually + # sum power spectrum functions and computes effective impedance + + # sum power functions + # BE CAREFUL: maybe for this "normalization sum" the sum should run + # on all single-bunch harmonics instead of only coupled-bunch + # harmonics (and then the frequency shift should be multiplied by + # n_bunches). This has to be checked. + sum1 = hmm_weighted_sum(m, nx) + + # effective impedance + sum2 = hmm_weighted_sum(m, nx, weight_function=impedance_function) + + effective_impedance[inx, im] = sum2 / sum1 + freq_shift = 1j * charge * single_bunch_current / (2 * (np.abs(m) + 1.) * m0 * gamma * tune * omega_rev * + bunch_length_seconds_meters) * sum2 / sum1 + + tune_shift = (freq_shift / omega_rev + m * omegas / omega_rev) + + tune_shift_nx[inx, im] = tune_shift + + # find the most unstable coupled-bunch mode for m=0 + inx = np.argmin(np.imag(tune_shift_nx[:, m_max])) + tune_shift_m0 = tune_shift_nx[inx, m_max] + + return tune_shift_nx, tune_shift_m0, effective_impedance diff --git a/xwakes/wit/utilities.py b/xwakes/wit/utilities.py new file mode 100644 index 00000000..148d5d41 --- /dev/null +++ b/xwakes/wit/utilities.py @@ -0,0 +1,732 @@ +from .component import Component +from .element import Element +from .interface import Layer, FlatIW2DInput, RoundIW2DInput +from .interface import component_names + +from yaml import load, SafeLoader +from typing import Tuple, Dict, List, Union, Sequence, Optional, Callable +from collections import defaultdict + +from numpy import (vectorize, sqrt, exp, pi, sin, cos, abs, sign, + inf, floor, linspace, ones, isscalar, array) +from numpy.typing import ArrayLike +import scipy.constants +from scipy import special as sp +import numpy as np + +if hasattr(np, 'trapezoid'): + trapz = np.trapezoid # numpy 2.0 +else: + trapz = np.trapz + +c_light = scipy.constants.speed_of_light # m s-1 +mu0 = scipy.constants.mu_0 +Z0 = mu0 * c_light +eps0 = scipy.constants.epsilon_0 + + +def string_to_params(name: str, include_is_impedance: bool = True): + """ + Converts a string describing some specific component to the set of parameters which give this component + :param name: The string description of the component to be converted. Either of the form where p is either + 'x', 'y' or 'z', and abcd are four integers giving the source- and test exponents of the component. Optionally, + a single character, 'z' or 'w' can be included at the front of the string if it is necessary to indicate that the + given component is either a wake- or impedance component. In this case, the include_is_impedance flag must be set. + :param include_is_impedance: A flag indicating whether or not "name" includes a 'z' or 'w' at the beginning. + :return: A tuple containing the necessary information to describe the type of the given component + """ + is_impedance = False + if include_is_impedance: + is_impedance = name[0] == 'z' + name = name[1:] + + plane = name[0].lower() + exponents = tuple(int(n) for n in name[1:]) + + return (is_impedance, plane, exponents) if include_is_impedance else (plane, exponents) + + +def create_component_from_config(identifier: str) -> Component: + """ + Creates a Component object as specified by a .yaml config file + :param identifier: A unique identifier-string corresponding to a component specification in config/component.yaml + :return: A Component object initialized according to the specification + """ + with open("config/component.yaml", "r") as f: + cdict = load(f, Loader=SafeLoader)[identifier] + wake = vectorize(lambda x: eval(cdict['wake'], {'exp': exp, 'sqrt': sqrt, + 'pi': pi, 'x': x})) if 'wake' in cdict else None + impedance = vectorize(lambda x: eval(cdict['impedance'], {'exp': exp, 'sqrt': sqrt, + 'pi': pi, 'x': x})) if 'impedance' in cdict else None + name = cdict['name'] if 'name' in cdict else "" + plane = cdict['plane'] + source_exponents = (int(cdict['source_exponents'][0]), int(cdict['source_exponents'][-1])) + test_exponents = (int(cdict['test_exponents'][0]), int(cdict['test_exponents'][-1])) + return Component(impedance, wake, plane, source_exponents, test_exponents, name) + + +def create_element_from_config(identifier: str) -> Element: + """ + Creates an Element object as specified by a .yaml config file + :param identifier: A unique identifier-string corresponding to an element specification in config/element.yaml + :return: An Element object initialized according to the specifications + """ + with open("config/element.yaml", "r") as f: + edict = load(f, Loader=SafeLoader)[identifier] + length = float(edict['length']) if 'length' in edict else 0 + beta_x = float(edict['beta_x']) if 'beta_x' in edict else 0 + beta_y = float(edict['beta_y']) if 'beta_y' in edict else 0 + name = edict['name'] if 'name' in edict else "" + tag = edict['tag'] if 'tag' in edict else "" + components = [create_component_from_config(c_id) for c_id in edict['components'].split()] \ + if 'components' in edict else [] + return Element(length, beta_x, beta_y, components, name, tag) + + +def compute_resonator_f_roi_half_width(q: float, f_r: float, f_roi_level: float = 0.5): + aux = np.sqrt((1 - f_roi_level) / f_roi_level) + + return (aux + np.sqrt(aux**2 + 4*q**2))*f_r/(2*q) - f_r + + +def create_resonator_component(plane: str, exponents: Tuple[int, int, int, int], + r: float, q: float, f_r: float, f_roi_level: float = 0.5) -> Component: + """ + Creates a single component object belonging to a resonator + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param r: the shunt impedance of the given component of the resonator + :param q: the quality factor of the given component of the resonator + :param f_r: the resonance frequency of the given component of the resonator + :param f_roi_level: fraction of the peak ok the resonator which is covered by the ROI. I.e. the roi will cover + the frequencies for which the resonator impedance is larger than f_roi_level*r + :return: A component object of a resonator, specified by the input arguments + """ + root_term = sqrt(1 - 1 / (4 * q ** 2) + 0J) + omega_r = 2 * pi * f_r + if plane == 'z': + impedance = lambda f: r / (1 - 1j * q * (f_r / f - f / f_r)) + omega_bar = omega_r * root_term + alpha = omega_r / (2 * q) + wake = lambda t: (omega_r * r * exp(-alpha * t) * ( + cos(omega_bar * t) - alpha * sin(omega_bar * t) / omega_bar) / q).real + else: + impedance = lambda f: (f_r * r) / (f * (1 - 1j * q * (f_r / f - f / f_r))) + wake = lambda t: (omega_r * r * exp(-omega_r * t / (2 * q)) * sin(omega_r * root_term * t) / + (q * root_term)).real + + if q > 1: + d = compute_resonator_f_roi_half_width(q=q, f_r=f_r, f_roi_level=f_roi_level) + f_rois = [(f_r - d, f_r + d)] + else: + f_rois = [] + # TODO: add ROI(s) for wake + + return Component(impedance, wake, plane, source_exponents=exponents[:2], + test_exponents=exponents[2:], + f_rois=f_rois) + + +def create_resonator_element(length: float, beta_x: float, beta_y: float, + rs: Dict[str, float], qs: Dict[str, float], fs: Dict[str, float], + f_roi_levels: Dict[str, float] = None, tag: str = 'resonator', + description: str = '') -> Element: + """ + Creates an element object representing a resonator. + :param length: The length, in meters, of the resonator element + :param beta_x: The value of the beta function in the x-direction at the position of the resonator element + :param beta_y: The value of the beta function in the y-direction at the position of the resonator element + :param rs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the + values give the Shunt impedance corresponding to this particular component + :param qs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the + values give the quality factor of the specified component of the resonator + :param fs: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and the + values give the resonance frequency corresponding to the particular component + :param f_roi_levels: A dictionary where the keys correspond to a plane followed by four exponents, i.e. "y0100", and + the values give the fraction of the peak ok the resonator which is covered by the ROI. I.e. the roi will cover + the frequencies for which the resonator impedance is larger than f_roi_level*r + :param tag: An optional short string used to place elements into categories + :param description: An optional short description of the element + :return: An element object as specified by the user-input + """ + if f_roi_levels is None: + f_roi_levels = {} + for key in rs.keys(): + f_roi_levels[key] = 0.5 + + assert set(rs.keys()) == set(qs.keys()) == set(fs.keys()) == set(f_roi_levels.keys()), "The three input " \ + "dictionaries describing " \ + "the resonator do not all " \ + "have identical keys" + components = [] + for key in rs.keys(): + plane, exponents = string_to_params(key, include_is_impedance=False) + components.append(create_resonator_component(plane, exponents, rs[key], qs[key], fs[key], + f_roi_level=f_roi_levels[key])) + + return Element(length, beta_x, beta_y, components, tag=tag, description=description) + + +def create_many_resonators_element(length: float, beta_x: float, beta_y: float, + params_dict: Dict[str, List[Dict[str, float]]], tag: str = 'resonator', + description: str = '') -> Element: + """ + Creates an element object representing many resonators. + :param length: The length, in meters, of the element + :param beta_x: The value of the beta function in the x-direction at the position of the element + :param beta_y: The value of the beta function in the y-direction at the position of the element + :param params_dict: a dictionary associating to each component a list of dictionaries containing the + parameters of a resonator component. E.g.: + params_dict = { + 'z0000': + [ + {'r': 10, 'q': 10, 'f': 50, 'f_roi_level: 0.5}, + {'r': 40, 'q': 100, 'f': 60} + ], + 'x1000': + ... + } + f_roi_level is optional + :param tag: An optional short string used to place elements into categories + :param description: An optional short description of the element + :return: An element object as specified by the user-input + """ + all_components = [] + for component_id, component_params_list in params_dict.items(): + plane, exponents = string_to_params(component_id, include_is_impedance=False) + for component_params in component_params_list: + assert ('r' in component_params.keys() and 'q' in component_params.keys() and + 'f' in component_params.keys()), "each of the the component dictionaries must contain r, q and f" + + f_roi_level = component_params.get('f_roi_level', 0.5) + all_components.append(create_resonator_component(plane, exponents, component_params['r'], + component_params['q'], component_params['f'], + f_roi_level=f_roi_level)) + + comp_dict = defaultdict(lambda: 0) + for c in all_components: + comp_dict[(c.plane, c.source_exponents, c.test_exponents)] += c + + components = sorted(comp_dict.values(), key=lambda x: (x.plane, x.source_exponents, x.test_exponents)) + + return Element(length, beta_x, beta_y, components, tag=tag, description=description) + + +def create_classic_thick_wall_component(plane: str, exponents: Tuple[int, int, int, int], + layer: Layer, radius: float) -> Component: + """ + Creates a single component object modeling a resistive wall impedance/wake, + based on the "classic thick wall formula" (see e.g. A. W. Chao, chap. 2 in + "Physics of Collective Beams Instabilities in High Energy Accelerators", + John Wiley and Sons, 1993). + Only longitudinal and transverse dipolar impedances are supported here. + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param layer: the chamber material, as a wit Layer object + :param radius: the chamber radius in m + :return: A component object + """ + + # Material properties required for the skin depth computation are derived from the input Layer attributes + material_resistivity = layer.dc_resistivity + material_relative_permeability = 1. + layer.magnetic_susceptibility + material_permeability = material_relative_permeability * mu0 + + # Create the skin depth as a function offrequency and layer properties + delta_skin = lambda f: (material_resistivity / (2*pi*abs(f) * material_permeability)) ** (1/2) + + # Longitudinal impedance and wake + if plane == 'z' and exponents == (0, 0, 0, 0): + impedance = lambda f: (1/2) * (1+sign(f)*1j) * material_resistivity / (pi * radius) * (1 / delta_skin(f)) + wake = lambda t: - c_light / (2*pi*radius) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(1/2)) + # Transverse dipolar impedance + elif (plane == 'x' and exponents == (1, 0, 0, 0)) or (plane == 'y' and exponents == (0, 1, 0, 0)): + impedance = lambda f: ((c_light/(2*pi*f)) * (1+sign(f)*1j) * + material_resistivity / (pi * radius**3) * + (1 / delta_skin(f))) + wake = lambda t: -c_light / (2*pi*radius**3) * (Z0 * material_resistivity/pi)**(1/2) * 1/(t**(3/2)) + else: + print("Warning: resistive wall impedance not implemented for component {}{}. Set to zero".format(plane, + exponents)) + impedance = lambda f: 0 + wake = lambda f: 0 + + return Component(vectorize(impedance), vectorize(wake), plane, source_exponents=exponents[:2], + test_exponents=exponents[2:]) + + +def _zlong_round_single_layer_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius: float, length: float) -> ArrayLike: + """ + Function to compute the longitudinal resistive-wall impedance from + the single-layer, approximated formula for a cylindrical structure, + by E. Metral (see e.g. Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, + https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius: the radius of the structure (in m) + :param length: the total length of the resistive object (in m) + :return: Zlong, the longitudinal impedance at these frequencies + """ + beta = sqrt(1.-1./gamma**2) + omega = 2*pi*frequencies + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*sqrt(1.-beta**2*eps1*mu1) + + coef_long = 1j*omega*mu0*length/(2.*pi*beta**2*gamma**2) + + x1 = k*radius/gamma + x1sq = x1**2 + x2 = nu*radius + + zlong = coef_long * (sp.k0(x1)/sp.i0(x1) - 1./(x1sq*(1./2.+eps1*sp.kve(1, x2)/(x2*sp.kve(0, x2))))) + + return zlong + + +def _zdip_round_single_layer_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius: float, length: float) -> ArrayLike: + """ + Function to compute the transverse dipolar resistive-wall impedance from + the single-layer, approximated formula for a cylindrical structure, + Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, + https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius: the radius of the structure (in m) + :param length: the total length of the resistive object (in m) + :return: Zdip, the transverse dipolar impedance at these frequencies + """ + beta = sqrt(1.-1./gamma**2) + omega = 2*pi*frequencies + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*sqrt(1.-beta**2*eps1*mu1) + + coef_dip = 1j*k**2*Z0*length/(4.*pi*beta*gamma**4) + + x1 = k*radius/gamma + x1sq = x1**2 + x2 = nu*radius + + zdip = coef_dip * (sp.k1(x1)/sp.i1(x1) + 4.*beta**2*gamma**2/(x1sq*(2.+x2*sp.kve(0, x2)/(mu1*sp.kve(1, x2))))) + + return zdip + + +def create_resistive_wall_single_layer_approx_component(plane: str, exponents: Tuple[int, int, int, int], + input_data: Union[FlatIW2DInput, RoundIW2DInput]) -> Component: + """ + Creates a single component object modeling a resistive wall impedance, + based on the single-layer approximated formulas by E. Metral (see e.g. + Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, + https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula the + Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle round + but the Yokoya factors put in the input will be used. + :return: A component object + """ + gamma = input_data.relativistic_gamma + length = input_data.length + + if isinstance(input_data, FlatIW2DInput): + if len(input_data.top_layers) > 1: + raise NotImplementedError("Input data can have only one layer") + yok_long = 1. + layer = input_data.top_layers[0] + radius = input_data.top_half_gap + if input_data.top_bottom_symmetry: + yok_dipx = pi**2/24. + yok_dipy = pi**2/12. + yok_quax = -pi**2/24. + yok_quay = pi**2/24. + elif input_data.bottom_half_gap == inf: + yok_dipx = 0.25 + yok_dipy = 0.25 + yok_quax = -0.25 + yok_quay = 0.25 + else: + raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " + "hence the bottom half gap must be infinite") + elif isinstance(input_data, RoundIW2DInput): + radius = input_data.inner_layer_radius + if len(input_data.layers) > 1: + raise NotImplementedError("Input data can have only one layer") + layer = input_data.layers[0] + yok_long = input_data.yokoya_factors[0] + yok_dipx = input_data.yokoya_factors[1] + yok_dipy = input_data.yokoya_factors[2] + yok_quax = input_data.yokoya_factors[3] + yok_quay = input_data.yokoya_factors[4] + else: + raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") + + # Longitudinal impedance + if plane == 'z' and exponents == (0, 0, 0, 0): + impedance = lambda f: yok_long*_zlong_round_single_layer_approx( + f, gamma, layer, radius, length) + # Transverse impedances + elif plane == 'x' and exponents == (1, 0, 0, 0): + impedance = lambda f: yok_dipx*_zdip_round_single_layer_approx( + f, gamma, layer, radius, length) + elif plane == 'y' and exponents == (0, 1, 0, 0): + impedance = lambda f: yok_dipy*_zdip_round_single_layer_approx( + f, gamma, layer, radius, length) + elif plane == 'x' and exponents == (0, 0, 1, 0): + impedance = lambda f: yok_quax*_zdip_round_single_layer_approx( + f, gamma, layer, radius, length) + elif plane == 'y' and exponents == (0, 0, 0, 1): + impedance = lambda f: yok_quay*_zdip_round_single_layer_approx( + f, gamma, layer, radius, length) + else: + impedance = lambda f: 0 + + return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], + test_exponents=exponents[2:]) + + +def create_resistive_wall_single_layer_approx_element( + input_data: Union[FlatIW2DInput, RoundIW2DInput], + beta_x: float, beta_y: float, + component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), + name: str = "", tag: str = "", description: str = "") -> Element: + """ + Creates an element object modeling a resistive wall impedance, + based on the single-layer approximated formulas by E. Metral (see e.g. + Eqs. 13-14 in N. Mounet and E. Metral, IPAC'10, TUPD053, + https://accelconf.web.cern.ch/IPAC10/papers/tupd053.pdf, and + Eq. 21 in F. Roncarolo et al, Phys. Rev. ST Accel. Beams 12, 084401, 2009, + https://doi.org/10.1103/PhysRevSTAB.12.084401) + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula the + Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle round + but the Yokoya factors put in the input will be used. + :param beta_x: The beta function in the x-plane at the position of the element + :param beta_y: The beta function in the y-plane at the position of the element + :param component_ids: a list of components to be computed + :param name: A user-specified name for the Element + :param tag: A string to tag the Element + :param description: A description for the Element + :return: An Element object representing the structure + """ + components = [] + length = input_data.length + for component_id in component_ids: + _, plane, exponents = component_names[component_id] + components.append(create_resistive_wall_single_layer_approx_component( + plane, exponents, input_data)) + + return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, + description=description) + + +def _zlong_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius_small: float, + radius_large: float, length: float, + step_size: float = 1e-3) -> ArrayLike: + """ + Function to compute the longitudinal resistive-wall impedance for a + round taper, integrating the radius-dependent approximated formula + for a cylindrical structure (see_zlong_round_single_layer_approx above), + over the length of the taper. + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param length: the total length of the taper (in m) + :param step_size: the step size (in the radial direction) for the + integration (in m) + :return: the longitudinal impedance at these frequencies + """ + if isscalar(frequencies): + frequencies = array(frequencies) + beta = sqrt(1.-1./gamma**2) + omega = 2*pi*frequencies.reshape((-1, 1)) + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*sqrt(1.-beta**2*eps1*mu1) + + coef_long = 1j*omega*mu0/(2.*pi*beta**2*gamma**2) + + npts = int(floor(abs(radius_large-radius_small)/step_size)+1) + radii = linspace(radius_small, radius_large, npts).reshape((1, -1)) + one_array = ones(radii.shape) + + x1 = k.dot(radii)/gamma + x1sq = x1**2 + x2 = nu.dot(radii) + zlong = (coef_long.dot(length / float(npts) * one_array) * + (sp.k0(x1) / sp.i0(x1) - 1. / (x1sq * (1. / 2. + eps1.dot(one_array) * + sp.kve(1, x2) / (x2 * sp.kve(0, x2))))) + ) + + return trapz(zlong, axis=1) + + +def _zdip_round_taper_RW_approx(frequencies: ArrayLike, gamma: float, + layer: Layer, radius_small: float, + radius_large: float, length: float, + step_size: float = 1e-3) -> ArrayLike: + """ + Function to compute the transverse dip. resistive-wall impedance for a + round taper, integrating the radius-dependent approximated formula + for a cylindrical structure (see_zdip_round_single_layer_approx above), + over the length of the taper. + :param frequencies: the frequencies (array) (in Hz) + :param gamma: relativistic mass factor + :param layer: a layer with material properties (only resistivity, + relaxation time and magnetic susceptibility are taken into account + at this stage) + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param length: the total length of the taper (in m) + :param step_size: the step size (in the radial direction) for the + integration (in m) + :return: the transverse dipolar impedance at these frequencies + """ + if isscalar(frequencies): + frequencies = array(frequencies) + beta = sqrt(1.-1./gamma**2) + omega = 2*pi*frequencies.reshape((-1,1)) + k = omega/(beta*c_light) + + rho = layer.dc_resistivity + tau = layer.resistivity_relaxation_time + mu1 = 1.+layer.magnetic_susceptibility + eps1 = 1. - 1j/(eps0*rho*omega*(1.+1j*omega*tau)) + nu = k*sqrt(1.-beta**2*eps1*mu1) + + coef_dip = 1j*k**2*Z0/(4.*pi*beta*gamma**4) + + npts = int(floor(abs(radius_large-radius_small)/step_size)+1) + radii = linspace(radius_small,radius_large,npts).reshape((1,-1)) + one_array = ones(radii.shape) + + x1 = k.dot(radii)/gamma + x1sq = x1**2 + x2 = nu.dot(radii) + zdip = ( + coef_dip.dot(length / float(npts) * one_array) * + (sp.k1(x1) / sp.i1(x1) + 4 * beta**2 * gamma**2 / (x1sq * (2 + x2 * sp.kve(0, x2) / (mu1 * sp.kve(1, x2))))) + ) + + return trapz(zdip, axis=1) + + +def create_taper_RW_approx_component(plane: str, exponents: Tuple[int, int, int, int], + input_data: Union[FlatIW2DInput, RoundIW2DInput], + radius_small: float, radius_large: float, + step_size: float = 1e-3) -> Component: + """ + Creates a single component object modeling a round or flat taper (flatness + along the horizontal direction, change of half-gap along the vertical one) + resistive-wall impedance, using the integration of the radius-dependent + approximated formula for a cylindrical structure (see + the above functions), over the length of the taper. + :param plane: the plane the component corresponds to + :param exponents: four integers corresponding to (source_x, source_y, test_x, test_y) aka (a, b, c, d) + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula the + Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle + round but the Yokoya factors put in the input will be used. + Note that the radius or half-gaps in input_data are not used (replaced + by the scan from radius_small to radius_large, for the integration). + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param step_size: the step size (in the radial or vertical direction) + for the integration (in m) + :return: A component object + """ + gamma = input_data.relativistic_gamma + length = input_data.length + + if isinstance(input_data, FlatIW2DInput): + if len(input_data.top_layers) > 1: + raise NotImplementedError("Input data can have only one layer") + yok_long = 1. + layer = input_data.top_layers[0] + radius = input_data.top_half_gap + if input_data.top_bottom_symmetry: + yok_dipx = pi**2/24. + yok_dipy = pi**2/12. + yok_quax = -pi**2/24. + yok_quay = pi**2/24. + elif input_data.bottom_half_gap == inf: + yok_dipx = 0.25 + yok_dipy = 0.25 + yok_quax = -0.25 + yok_quay = 0.25 + else: + raise NotImplementedError("For asymmetric structures, only the case of a single plate is implemented; " + "hence the bottom half gap must be infinite") + elif isinstance(input_data, RoundIW2DInput): + radius = input_data.inner_layer_radius + if len(input_data.layers) > 1: + raise NotImplementedError("Input data can have only one layer") + layer = input_data.layers[0] + yok_long = input_data.yokoya_factors[0] + yok_dipx = input_data.yokoya_factors[1] + yok_dipy = input_data.yokoya_factors[2] + yok_quax = input_data.yokoya_factors[3] + yok_quay = input_data.yokoya_factors[4] + else: + raise NotImplementedError("Input of type neither FlatIW2DInput nor RoundIW2DInput cannot be handled") + + # Longitudinal impedance + if plane == 'z' and exponents == (0, 0, 0, 0): + impedance = lambda f: yok_long*_zlong_round_taper_RW_approx( + f, gamma, layer, radius_small, radius_large, + length, step_size=step_size) + # Transverse impedances + elif plane == 'x' and exponents == (1, 0, 0, 0): + impedance = lambda f: yok_dipx*_zdip_round_taper_RW_approx( + f, gamma, layer, radius_small, radius_large, + length, step_size=step_size) + elif plane == 'y' and exponents == (0, 1, 0, 0): + impedance = lambda f: yok_dipy*_zdip_round_taper_RW_approx( + f, gamma, layer, radius_small, radius_large, + length, step_size=step_size) + elif plane == 'x' and exponents == (0, 0, 1, 0): + impedance = lambda f: yok_quax*_zdip_round_taper_RW_approx( + f, gamma, layer, radius_small, radius_large, + length, step_size=step_size) + elif plane == 'y' and exponents == (0, 0, 0, 1): + impedance = lambda f: yok_quay*_zdip_round_taper_RW_approx( + f, gamma, layer, radius_small, radius_large, + length, step_size=step_size) + else: + impedance = lambda f: 0 + + return Component(impedance=impedance, plane=plane, source_exponents=exponents[:2], + test_exponents=exponents[2:]) + + +def create_taper_RW_approx_element( + input_data: Union[FlatIW2DInput, RoundIW2DInput], + beta_x: float, beta_y: float, + radius_small: float, radius_large: float, step_size: float=1e-3, + component_ids: Sequence[str] = ('zlong', 'zxdip', 'zydip', 'zxqua', 'zyqua'), + name: str = "", tag: str = "", description: str = "") -> Element: + """ + Creates an element object modeling a round or flat taper (flatness + along the horizontal direction, change of half-gap along the vertical one) + resistive-wall impedance, using the integration of the radius-dependent + approximated formula for a cylindrical structure (see + the above functions), over the length of the taper. + :param input_data: an IW2D input object (flat or round). If the input + is of type FlatIW2DInput and symmetric, we apply to the round formula the + Yokoya factors for an infinitely flat structure (see e.g. K. Yokoya, + KEK Preprint 92-196 (1993), and Part. Accel. 41 (1993) pp.221-248, + https://cds.cern.ch/record/248630/files/p221.pdf), + while for a single plate we use those from A. Burov and V. Danilov, + PRL 82,11 (1999), https://doi.org/10.1103/PhysRevLett.82.2286. Other + kinds of asymmetric structure will raise an error. + If the input is of type RoundIW2DInput, the structure is in principle round + but the Yokoya factors put in the input will be used. + Note that the radius or half-gaps in input_data are not used (replaced + by the scan from radius_small to radius_large, for the integration). + :param beta_x: The beta function in the x-plane at the position of the element + :param beta_y: The beta function in the y-plane at the position of the element + :param radius_small: the smallest radius of the taper (in m) + :param radius_large: the largest radius of the taper (in m) + :param step_size: the step size (in the radial or vertical direction) + for the integration (in m) + :param component_ids: a list of components to be computed + :param name: A user-specified name for the Element + :param tag: A string to tag the Element + :param description: A description for the Element + :return: An Element object representing the structure + """ + components = [] + length = input_data.length + for component_id in component_ids: + _, plane, exponents = component_names[component_id] + components.append(create_taper_RW_approx_component(plane=plane, exponents=exponents, input_data=input_data, + radius_small=radius_small, radius_large=radius_large, + step_size=step_size)) + + return Element(length=length, beta_x=beta_x, beta_y=beta_y, components=components, name=name, tag=tag, + description=description) + + +def create_interpolated_impedance_component(interpolation_frequencies: ArrayLike, + impedance: Optional[Callable] = None, + wake: Optional[Callable] = None, plane: str = '', + source_exponents: Tuple[int, int] = (-1, -1), + test_exponents: Tuple[int, int] = (-1, -1), + name: str = "Unnamed Component", + f_rois: Optional[List[Tuple[float, float]]] = None, + t_rois: Optional[List[Tuple[float, float]]] = None): + """ + Creates a component in which the impedance function is evaluated directly only on few points and it is interpolated + everywhere else. This helps when the impedance function is very slow to evaluate. + :param interpolation_frequencies: the frequencies where the function is evaluated for the interpolation + :param impedance: A callable function representing the impedance function of the Component. Can be undefined if + the wake function is defined. + :param wake: A callable function representing the wake function of the Component. Can be undefined if + the impedance function is defined. + :param plane: The plane of the Component, either 'x', 'y' or 'z'. Must be specified for valid initialization + :param source_exponents: The exponents in the x and y planes experienced by the source particle. Also + referred to as 'a' and 'b'. Must be specified for valid initialization + :param test_exponents: The exponents in the x and y planes experienced by the source particle. Also + referred to as 'a' and 'b'. Must be specified for valid initialization + :param name: An optional user-specified name of the component + :param f_rois: a list of frequency regions of interest + :param t_rois: a list of time regions of interest + """ + def interpolated_impedance(x): + return np.interp(x, interpolation_frequencies, impedance(interpolation_frequencies)) + + return Component(impedance=interpolated_impedance, wake=wake, plane=plane, + source_exponents=source_exponents, test_exponents=test_exponents, + name=name, f_rois=f_rois, + t_rois=t_rois) diff --git a/xwakes/wit/utils.py b/xwakes/wit/utils.py new file mode 100644 index 00000000..cff25e4b --- /dev/null +++ b/xwakes/wit/utils.py @@ -0,0 +1,35 @@ +from typing import Sequence, Union + +import numpy as np + + +def round_sigfigs(arr: Union[float, Sequence[float]], sigfigs: int): + """ + Rounds all the floats in an array to a the given number of + significant digits. + Warning: this routine is slow when working with very long + arrays. If the final goal is to get unique elements to the given + precision, it is better to use unique_sigfigs. + :param arr: A float or an array-like sequence of floats. + :param sigfigs: The number of significant digits (integer). + :return: The same array (or float) rounded appropriately. + """ + if np.isscalar(arr): + return np.round(arr, sigfigs - 1 - int(np.floor(np.log10(np.abs(arr))))) + else: + return np.array([round_sigfigs(value, sigfigs) for value in arr]) + + +def unique_sigfigs(arr: np.ndarray, sigfigs: int): + """ + Given an array, returns it sorted and with only the elements that + are unique after rounding to a given number of significant digits. + :param arr: An array of floats + :param sigfigs: The number of significant digits (integer) + :return: the sorted array with unique elements + """ + arr = np.unique(arr.flat) + sorted_indices = np.argsort(arr.flat) + log_diff = np.floor(np.log10(np.append(np.inf, np.diff(arr[sorted_indices])))) + log_arr = np.floor(np.log10(arr[sorted_indices])) + return arr.flat[sorted_indices[log_diff - log_arr >= -sigfigs]]