From 504f40edd2e73deaea98c0da08ccf42390e98080 Mon Sep 17 00:00:00 2001 From: Benoit Sierro Date: Thu, 23 Mar 2023 09:39:51 +0100 Subject: [PATCH] some slight cleanup --- src/scgenerator/math.py | 173 ++++++++++++------------------ src/scgenerator/physics/plasma.py | 9 +- 2 files changed, 70 insertions(+), 112 deletions(-) diff --git a/src/scgenerator/math.py b/src/scgenerator/math.py index fe0b414..bd451ef 100644 --- a/src/scgenerator/math.py +++ b/src/scgenerator/math.py @@ -19,25 +19,34 @@ def expm1_int(y: np.ndarray, dx: float) -> np.ndarray: return -np.expm1(-cumulative_simpson(y) * dx) -def span(*vec): +def span(*vec: np.ndarray) -> tuple[float, float]: """returns the min and max of whatever array-like is given. can accept many args""" out = (np.inf, -np.inf) + if len(vec) == 0 or len(vec[0]) == 0: + raise ValueError(f"did not provide any value to span") for x in vec: x = np.atleast_1d(x) - if len(x.shape) > 1: - x = x.ravel() - minx = np.min(x) - maxx = np.max(x) - out = (np.min([minx, out[0]]), np.max([maxx, out[1]])) - - if out[0] == np.inf or out[1] == -np.inf: - out = (0, 1) - print(f"failed to evaluate the span of {vec}") + out = (np.min([np.min(x), out[0]]), np.max([np.max(x), out[1]])) return out -def argclosest(array: np.ndarray, target: Union[float, int]): - """returns the index/indices corresponding to the closest matches of target in array""" +def argclosest(array: np.ndarray, target: float | int | Sequence[float | int]) -> int | np.ndarray: + """ + returns the index/indices corresponding to the closest matches of target in array + + Parameters + ---------- + array : np.ndarray, shape (n,) + array of values + target : number | np.ndarray + find the closest value to target in `array`. The index of the closest match is returned. + + Returns + ------- + int | np.ndarray + index / indices of the closest match + + """ min_dist = np.inf index = None if isinstance(target, (list, tuple, np.ndarray)): @@ -51,7 +60,7 @@ def argclosest(array: np.ndarray, target: Union[float, int]): return index -def length(x): +def extent(x): return np.max(x) - np.min(x) @@ -109,85 +118,6 @@ def u_nm(n, m): return jn_zeros(n - 1, m)[-1] -@np_cache -def ndft_matrix(t: np.ndarray, f: np.ndarray) -> np.ndarray: - """creates the nfft matrix - - Parameters - ---------- - t : np.ndarray, shape = (n,) - time array - f : np.ndarray, shape = (m,) - frequency array - - Returns - ------- - np.ndarray, shape = (m, n) - multiply x(t) by this matrix to get ~X(f) - """ - P, F = np.meshgrid(t, f) - return np.exp(-2j * np.pi * P * F) - - -@np_cache -def indft_matrix(t: np.ndarray, f: np.ndarray) -> np.ndarray: - """creates the nfft matrix - - Parameters - ---------- - t : np.ndarray, shape = (n,) - time array - f : np.ndarray, shape = (m,) - frequency array - - Returns - ------- - np.ndarray, shape = (m, n) - multiply ~X(f) by this matrix to get x(t) - """ - return np.linalg.pinv(ndft_matrix(t, f)) - - -def ndft(t: np.ndarray, s: np.ndarray, f: np.ndarray) -> np.ndarray: - """computes the Fourier transform of an uneven signal - - Parameters - ---------- - t : np.ndarray, shape = (n,) - time array - s : np.ndarray, shape = (n, ) - amplitute at each point of t - f : np.ndarray, shape = (m, ) - desired frequencies - - Returns - ------- - np.ndarray, shape = (m, ) - amplitude at each frequency - """ - return ndft_matrix(t, f) @ s - - -def indft(f: np.ndarray, a: np.ndarray, t: np.ndarray) -> np.ndarray: - """computes the inverse Fourier transform of an uneven spectrum - - Parameters - ---------- - f : np.ndarray, shape = (n,) - frequency array - a : np.ndarray, shape = (n, ) - amplitude at each point of f - t : np.ndarray, shape = (m, ) - time array - - Returns - ------- - np.ndarray, shape = (m, ) - amplitude at each point of t - """ - return indft_matrix(t, f) @ a - - def all_zeros(x: np.ndarray, y: np.ndarray) -> np.ndarray: """find all the x values such that y(x)=0 with linear interpolation""" pos = np.argwhere(y[1:] * y[:-1] < 0)[:, 0] @@ -196,7 +126,9 @@ def all_zeros(x: np.ndarray, y: np.ndarray) -> np.ndarray: def wspace(t, t_num=0): - """frequency array such that x(t) <-> np.fft(x)(w) + """ + frequency array such that x(t) <-> np.fft(x)(w) + Parameters ---------- t : float or array @@ -204,8 +136,9 @@ def wspace(t, t_num=0): array : time array t_num : int- if t is a float, specifies the number of points + Returns - ---------- + ------- w : array linspace of frencies corresponding to t """ @@ -221,7 +154,9 @@ def wspace(t, t_num=0): def tspace(time_window=None, t_num=None, dt=None): - """returns a time array centered on 0 + """ + returns a time array centered on 0 + Parameters ---------- time_window : float @@ -257,7 +192,8 @@ def tspace(time_window=None, t_num=None, dt=None): def build_envelope_w_grid(t: np.ndarray, w0: float) -> tuple[np.ndarray, np.ndarray, np.ndarray]: - """computes a bunch of values that relate to the simulation grid + """ + convenience function to Parameters ---------- @@ -334,7 +270,36 @@ def build_t_grid( return t, time_window, dt, t_num -def polynom_extrapolation(x: np.ndarray, y: np.ndarray, degree: float) -> np.ndarray: +def polynom_extrapolation(x: np.ndarray, y: np.ndarray, degree: int) -> np.ndarray: + """ + performs a polynomial extrapolation on both ends of the support + + Parameters + ---------- + x : np.ndarray (n,) + x values + y : np.ndarray (n,) + y values. The shape must correspond to that of x. + degree : int + degree of the polynom. + + Returns + ------- + np.ndarray, shape (n,) + y array with zero values on either side replaces with polynomial extrapolation + + Example + + + + """ + out = y.copy() + order = np.argsort(x) + left_ind, *_, right_ind = np.nonzero(out[order])[0] + return _polynom_extrapolation_in_place(out[order], left_ind, right_ind, degree)[order.argsort] + + +def _polynom_extrapolation_in_place(y: np.ndarray, left_ind: int, right_ind: int, degree: float): """extrapolates IN PLACE linearily on both side of the support Parameters @@ -346,13 +311,6 @@ def polynom_extrapolation(x: np.ndarray, y: np.ndarray, degree: float) -> np.nda right_ind : int last value we want to keep (extrapolate to the right of that) """ - out = y.copy() - order = np.argsort(x) - left_ind, *_, right_ind = np.nonzero(out[order])[0] - return _polynom_extrapolation_in_place(out[order], left_ind, right_ind, degree)[order.argsort] - - -def _polynom_extrapolation_in_place(y: np.ndarray, left_ind: int, right_ind: int, degree: float): r_left = (1 + np.arange(left_ind))[::-1] ** degree r_right = np.arange(len(y) - right_ind) ** degree y[:left_ind] = r_left * (y[left_ind] - y[left_ind + 1]) + y[left_ind] @@ -390,7 +348,8 @@ def envelope_ind( signal: np.ndarray, dmin: int = 1, dmax: int = 1, split: bool = False ) -> tuple[np.ndarray, np.ndarray]: """ - returns the indices of the top/bottom envelope of a signal + Attempts to separate the envolope from a periodic signal and return the indices of the + top/bottom envelope of a signal Parameters ---------- @@ -405,9 +364,9 @@ def envelope_ind( Returns ------- - np.ndarray, shape (m,), m < n + np.ndarray, shape (m,), m <= n lower envelope indices - np.ndarray, shape (l,), l < n + np.ndarray, shape (l,), l <= n upper envelope indices """ diff --git a/src/scgenerator/physics/plasma.py b/src/scgenerator/physics/plasma.py index b50ac53..68443f7 100644 --- a/src/scgenerator/physics/plasma.py +++ b/src/scgenerator/physics/plasma.py @@ -2,7 +2,6 @@ from typing import Any, Callable, NamedTuple import numpy as np import scipy.special -from numpy.core.numeric import zeros_like from scipy.interpolate import interp1d from scgenerator.math import cumulative_simpson, expm1_int @@ -94,14 +93,14 @@ class Plasma: number density of free electrons as function of time """ field_abs: np.ndarray = np.abs(field) - nzi = field != 0 - rate = zeros_like(field_abs) - rate[nzi] = self.rate(field_abs[nzi]) + valid_ind = field != 0 + rate = np.zeros_like(field_abs) + rate[valid_ind] = self.rate(field_abs[valid_ind]) electron_density = free_electron_density(rate, self.dt, N0) dn_dt: np.ndarray = (N0 - electron_density) * rate loss_term = np.zeros_like(field) - loss_term[nzi] = dn_dt[nzi] * self.ionization_energy / field[nzi] + loss_term[valid_ind] = dn_dt[valid_ind] * self.ionization_energy / field[valid_ind] phase_term = self.dt * e**2 / me * cumulative_simpson(electron_density * field)