Feature impedance to wake conversion

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.

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:
         """