From cd16a243cc85dc932accb3f490de1540f4595907 Mon Sep 17 00:00:00 2001 From: Eskil Vik Date: Wed, 26 Apr 2023 16:42:01 +0200 Subject: [PATCH 1/4] refactor f and t array generation --- pywit/component.py | 70 ++++++++++++++++++++++------------------------ 1 file changed, 34 insertions(+), 36 deletions(-) diff --git a/pywit/component.py b/pywit/component.py index a4314897..fd759a7a 100644 --- a/pywit/component.py +++ b/pywit/component.py @@ -10,7 +10,7 @@ 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 +24,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: @@ -99,7 +106,7 @@ 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: @@ -284,8 +291,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 @@ -305,23 +324,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)) - - xs = mix_fine_and_rough_sampling(start, stop, rough_points, - fine_points_per_roi, - f_rois_no_dup) + frequencies = self._frequency_array(rough_points, start, stop, precision_factor) + + return frequencies, self.impedance(frequencies) - 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 @@ -341,20 +350,9 @@ 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, - 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, From e03be9c615e6a7675653da144becf6ff927fa485 Mon Sep 17 00:00:00 2001 From: Eskil Vik Date: Fri, 28 Apr 2023 11:59:39 +0200 Subject: [PATCH 2/4] add draft for wake to impedance function --- pywit/component.py | 37 ++++++++++++++++++++++++++++--------- 1 file changed, 28 insertions(+), 9 deletions(-) diff --git a/pywit/component.py b/pywit/component.py index fd759a7a..073f5d02 100644 --- a/pywit/component.py +++ b/pywit/component.py @@ -5,6 +5,8 @@ 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 @@ -91,16 +93,33 @@ def generate_wake_from_impedance(self) -> None: 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(rough_points=1000) + frequencies = self._frequency_array(rough_points=5000) + + 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: """ From 2280dadb46f99903fa9552b2046436f09dadc453 Mon Sep 17 00:00:00 2001 From: Eskil Vik Date: Thu, 1 Jun 2023 10:15:23 +0200 Subject: [PATCH 3/4] add (quickly) points tuning in wake conversion --- pywit/component.py | 8 +++++--- 1 file changed, 5 insertions(+), 3 deletions(-) diff --git a/pywit/component.py b/pywit/component.py index 073f5d02..b4f37769 100644 --- a/pywit/component.py +++ b/pywit/component.py @@ -87,7 +87,9 @@ 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. @@ -101,8 +103,8 @@ def generate_wake_from_impedance(self) -> None: # 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(rough_points=1000) - frequencies = self._frequency_array(rough_points=5000) + 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, From 3dc748d944f47f80cb78aef28bb7ac140df55a98 Mon Sep 17 00:00:00 2001 From: Eskil Vik Date: Thu, 1 Jun 2023 10:35:41 +0200 Subject: [PATCH 4/4] add (draft) method in Element to calculate wakes --- pywit/element.py | 26 +++++++++++++++++++++++++- 1 file changed, 25 insertions(+), 1 deletion(-) diff --git a/pywit/element.py b/pywit/element.py index d12a9c05..e509ce24 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 @@ -155,6 +155,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: """