diff --git a/pywit/component.py b/pywit/component.py index 49818dce..d103987b 100644 --- a/pywit/component.py +++ b/pywit/component.py @@ -5,12 +5,14 @@ from typing import Optional, Callable, Tuple, Union, List +from neffint.fixed_grid_fourier_integral import fourier_integral_fixed_sampling +from scipy.interpolate import PchipInterpolator 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]]): + precision_factor: float, rois: List[Tuple[float, float]]): """ Mix a fine and rough (geometric) sampling between start and stop, refined in the regions of interest rois. @@ -24,14 +26,21 @@ def mix_fine_and_rough_sampling(start: float, stop: float, rough_points: int, 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)] + if len(rois) == 0: + return np.geomspace(start, stop, rough_points) + + # eliminate duplicates + rois_no_dup = set(rois) + + fine_points_per_roi = int(round(rough_points * precision_factor)) + + intervals = [np.linspace(max(roi_start, start), min(roi_stop, stop), fine_points_per_roi) + for roi_start, roi_stop in rois_no_dup + if (start <= roi_start <= stop or start <= roi_stop <= stop)] fine_sampling_rois = np.hstack(intervals) if intervals else np.array([]) rough_sampling = np.geomspace(start, stop, rough_points) - return np.sort(list(set(round_sigfigs( - snp.merge(fine_sampling_rois, rough_sampling), 7)))) + return np.unique(round_sigfigs(snp.merge(fine_sampling_rois,rough_sampling),7)) class Component: @@ -79,28 +88,47 @@ def __init__(self, impedance: Optional[Callable] = None, wake: Optional[Callable 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: + def generate_wake_from_impedance(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) -> 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 + # 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." + + # TODO: Don't use magic number 1000 and 5000 here, and maybe not defaults for the rest. + # Find a good place to have the user input them + times = self._time_array(time_points, time_start, time_stop, time_precision_factor) + frequencies = self._frequency_array(freq_points, freq_start, freq_stop, freq_precision_factor) + + transform_arr = fourier_integral_fixed_sampling( + times=times, + frequencies=frequencies, + func_values=self.impedance(frequencies), + inf_correction_term=True, + interpolation="pchip" + ) + + if self.plane == "z": + wake_arr = np.real(transform_arr)/np.pi + else: + wake_arr = np.imag(transform_arr)/np.pi + + self.wake = PchipInterpolator(times, wake_arr) + return - # 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 + :return: Nothing@ """ # # If the object already has an impedance function, there is no need to generate it. # if self.impedance: @@ -285,8 +313,20 @@ def __eq__(self, other: Component) -> bool: 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, + + def _time_array(self, rough_points: int, start: float = MIN_TIME, stop: float = MAX_TIME, + precision_factor: float = TIME_P_FACTOR) -> np.ndarray: + + return mix_fine_and_rough_sampling(start, stop, rough_points, precision_factor, self.t_rois) + + + def _frequency_array(self, rough_points: int, start: float = MIN_FREQ, stop: float = MAX_FREQ, + precision_factor: float = FREQ_P_FACTOR) -> np.ndarray: + + return mix_fine_and_rough_sampling(start, stop, rough_points, precision_factor, self.f_rois) + + + 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 @@ -306,23 +346,13 @@ def impedance_to_array(self, rough_points: int, start: float = MIN_FREQ, :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)) + + frequencies = self._frequency_array(rough_points, start, stop, precision_factor) + + return frequencies, self.impedance(frequencies) - 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, + 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 @@ -342,26 +372,15 @@ def wake_to_array(self, rough_points: int, start: float = MIN_TIME, :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)) + times = self._time_array(rough_points, start, stop, precision_factor) - return xs, self.wake(xs) + return times, self.wake(times) 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]]: + 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. diff --git a/pywit/element.py b/pywit/element.py index cc0c4322..80a09d8c 100644 --- a/pywit/element.py +++ b/pywit/element.py @@ -2,7 +2,7 @@ from pywit.component import Component, Union -from typing import List +from typing import List, Dict, Optional from collections import defaultdict from scipy.special import comb @@ -157,6 +157,30 @@ def changed_betas(self, new_beta_x: float, new_beta_y: float) -> Element: y_ratio = self.beta_y / new_beta_y new_components = [((x_ratio ** c.power_x) * (y_ratio ** c.power_y)) * c for c in self.components] return Element(self.length, new_beta_x, new_beta_y, new_components, self.name, self.tag, self.description) + + + def calculate_all_element_wakes(self, discretization_kwarg_list: Optional[List[Dict[str, Union[int, float]]]]=None) -> None: + """Generate wakes for all constituent components. + + :param discretization_kwarg_list: A list of the same length as `self.components` of dictionaries containing the keyword arguments for + the `Component.generate_wake_from_impedance` method. Each dictionary takes a string (the argument name) as a key and the value will be passed + as a parameter to the `generate_wake_from_impedance` method. See example below. None (default) if no keywords should be given + :type discretization_kwarg_list: Optional[List[Dict[str, Union[int, float]]]], optional + + Example: Make the 3rd (0-indexed) component calculate the wake with 10k points: + my_element = ... # Generate an Element + discretization_kwargs_list = [{} for component in my_element.components] + discretization_kwargs_list[3]["time_points"] = 10000 + """ + + if discretization_kwarg_list is None: + discretization_kwarg_list = [{} for component in self.components] + + assert len(discretization_kwarg_list) == len(self.components), "discretization_kwarg_list must be None or a list of the same length as self.components" + + for component, discretization_kwargs in zip(self.components, discretization_kwarg_list): + component.generate_wake_from_impedance(**discretization_kwargs) + def __add__(self, other: Element) -> Element: """