pypulse-unal 0.1.1

A Python library for pulse analysis.

PyPulse

This python library is a JAX-based tool for pulse construction, analysis, and visualization, using JAX for high-performance calculations and GPU acceleration. Pypulse provides JAX-compatible functions for constructing electric fields in the frequency domain, applying phase modulation (chirp) and amplitude filtering (notch), performing Fourier transforms to the time domain, and calculating advanced pulse representations like the Wigner function.

Installation

You can install the library directly from PyPI:

pip install pypulse_unal

---

Function Documentation

The functions rely on 2 modules:

• pulses: Provides pulse shape definitions and constants.
• aux_functions: Provides Fourier transform and numerical limit finding utilities.

frequency_field

def frequency_field(
    f: jnp.ndarray,
    f_0: float,
    pulse_fwhm: float,
    pulse_area: float,
    pulse_type: int
) -> jnp.ndarray:

Constructs a pure electric field pulse in the frequency domain based on a specified shape, bandwidth, and area.

Type	Name	Description
Input	f (jnp.ndarray)	Array of frequencies where the pulse is evaluated.
Input	f_0 (float)	Pulse center frequency (carrier frequency).
Input	pulse_fwhm (float)	Pulse bandwidth (Full Width at Half Maximum) in the power profile.
Input	pulse_area (float)	Scaling factor for the pulse magnitude, representing the pulse area.
Input	pulse_type (int)	An integer identifier defining the shape of the pulse (e.g., Gaussian, Sech$^2$).
Output	jnp.ndarray	The evaluated pulse in the frequency domain, $E(f)$.

Operation: Calculates $E(f)$ using a shape function $S(f)$ from the external pulses module: $$E(f) = \text{pulse_area} \cdot S\left(f, f_0, \frac{\text{pulse_fwhm}}{\text{const}}\right)$$ The shape function's width is normalized using the constant corresponding to pulse_type to ensure the bandwidth parameter truly represents the FWHM.

---

set_chirp

def set_chirp(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    f_0: float,
    chirp_alpha: float,
    chrip_power: int = 1
) -> jnp.ndarray:

Applies a spectral phase (chirp) to an existing frequency field. The phase is a polynomial function of the frequency offset $(\omega - \omega_0)$.

Type	Name	Description
Input	f (jnp.ndarray)	Array of frequencies.
Input	f_field (jnp.ndarray)	The un-chirped frequency field array, $E(f)$.
Input	f_0 (float)	Chirp center frequency (the frequency around which the phase is zero).
Input	chirp_alpha (float)	The alpha parameter ($\alpha$) controlling the strength/magnitude of the phase.
Input	chirp_power (int)	The power of the frequency offset $(\omega - \omega_0)$ in the phase term (e.g., $1$ for linear phase, $2$ for quadratic/GDD, $3$ for cubic/TOD).
Output	jnp.ndarray	The chirped pulse in the frequency domain, $E'(f)$.

Operation: Applies the phase factor $\exp(i \phi(f))$ to the field $E(f)$. The phase $\phi(f)$ is given by: $$\phi(f) = \alpha \cdot \frac{\left(2\pi(f-f_0)\right)^{\text{chirp_power} + 1}}{\text{chirp_power} + 1}$$ If $\text{chirp_power} = 0$, a constant phase $\phi(f) = \alpha$ is applied.

---

set_notch

def set_notch(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    f_0: float,
    notch_fwhm: float,
    notch_type: int = pulses.gaussian
) -> jnp.ndarray:

Applies a notch filter (frequency amplitude filter) to the pulse, effectively removing power within a certain frequency band.

Type	Name	Description
Input	f (jnp.ndarray)	Array of frequencies.
Input	f_field (jnp.ndarray)	The frequency field array, $E(f)$.
Input	f_0 (float)	Notch center frequency.
Input	notch_fwhm (float)	Notch bandwidth (FWHM of the removal function).
Input	notch_type (int)	An integer identifier defining the shape of the notch filter (e.g., Gaussian).
Output	jnp.ndarray	The filtered pulse in the frequency domain.

Operation: The pulse is multiplied by a filter function $F(f)$: $$E'(f) = E(f) \cdot F(f)$$ where $F(f) = 1 - S(f)$. The function $S(f)$ is the specified shape (notch_type) centered at $f_0$ with width $\text{notch_fwhm}$.

---

time_field

def time_field(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: jnp.ndarray,
    t_0: float = 0
) -> jnp.ndarray:

Calculates the pulse electric field in the temporal domain, $E(t)$, by performing the inverse Fourier transform of the frequency field $E(f)$.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array (angular frequencies $\omega = 2\pi f$).
Input	f_field (jnp.ndarray)	Frequency field array, $E(f)$.
Input	t (jnp.ndarray)	Time array where the pulse is evaluated.
Input	t_0 (float)	Center of the pulse in time (temporal shift).
Output	jnp.ndarray	The pulse in the temporal domain, $E(t)$.

Operation: Uses the external aux_functions.vectorized_inverse_fourier to compute: $$E(t - t_0) = \int E(f) e^{i 2\pi f (t - t_0)} df$$

---

set_time_limits

def set_time_limits(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    step: float,
    t_0: float = 0,
    eps: float = 1e-5,
    max_steps: int = 10_000,
    negative: bool = False
) -> float:

Determines the temporal boundary (limit) where the pulse's intensity drops below a specified minimum value ($\epsilon$). Useful for setting the boundaries of the time grid.

Type	Name	Description
Input	f (jnp.ndarray)	Frequency domain array.
Input	f_field (jnp.ndarray)	Pulse in frequency domain.
Input	step (float)	Size of the steps used for searching the limit.
Input	t_0 (float)	Center of the interval to search for the limit.
Input	eps (float)	Minimum normalized intensity value used to set the limit (e.g., $10^{-5}$).
Input	max_steps (float)	Maximum number of iterations allowed to find the limit.
Input	negative (bool)	If True, searches for the limit in the negative time direction.
Output	float	The temporal limit $t_{lim}$ where $

Operation: This function internally defines the time-domain pulse function, then calls the generic numerical search utility aux_functions.set_limits to iteratively find the time coordinate where the pulse intensity falls below eps.

---

set_frequency_limits

def set_frequency_limits(
    f_0: float,
    pulse_fwhm: float,
    pulse_type: int,
    step: float,
    eps: float = 1e-5,
    max_steps: int = 10_000,
    negative: bool = False
) -> float:

Determines the frequency boundary (limit) where the pulse's intensity drops below a specified minimum value ($\epsilon$). Useful for setting the boundaries of the frequency grid.

Type	Name	Description
Input	f_0 (float)	Center of the interval to search for the limit (pulse center frequency).
Input	pulse_fwhm (float)	Pulse bandwidth (used to normalize the pulse width).
Input	pulse_type (int)	Shape of the pulse in frequency domain.
Input	step (float)	Size of the steps used for searching the limit.
Input	eps (float)	Minimum normalized intensity value used to set the limit.
Input	max_steps (float)	Maximum number of iterations allowed to find the limit.
Input	negative (bool)	If True, searches for the limit in the negative frequency direction.
Output	float	The frequency limit $f_{lim}$ where $

Operation: This function internally defines the frequency-domain pulse intensity function, then calls the generic numerical search utility aux_functions.set_limits to iteratively find the frequency coordinate where the pulse intensity falls below eps.

---

get_inst_frequency

def get_inst_frequency(
        f: jnp.ndarray,
        f_field: jnp.ndarray,
        t: jnp.ndarray,
        t_0: float):

Calculates the instantaneous frequency $\omega_{inst}(t)$ of the pulse in the temporal domain.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array.
Input	f_field (jnp.ndarray)	Frequency field array, $E(f)$.
Input	t (jnp.ndarray)	Time array where the instantaneous frequency is evaluated.
Input	t_0 (float)	Temporal shift of the pulse.
Output	jnp.ndarray	The instantaneous frequency $\omega_{inst}(t)$.

Operation: Calculates the instantaneous frequency using the phase $\phi(t)$ of the temporal electric field $E(t) = |E(t)| e^{i\phi(t)}$: $$\omega_{inst}(t) = \frac{1}{2\pi} \frac{d\phi(t)}{dt}$$ This calculation is delegated to aux_functions.aux_instantaneous_frequency.

---

wigner_function

def wigner_function(
    t_plot: jnp.ndarray,
    f_plot: jnp.ndarray,
    f_field: jnp.ndarray,
) -> jnp.ndarray:

Calculates the Wigner function $W(t, f)$, a time-frequency distribution that visualizes the intensity and chirp of the pulse.

Type	Name	Description
Input	t_plot (jnp.ndarray)	Time grid for the Wigner function evaluation (first output dimension).
Input	f_plot (jnp.ndarray)	Frequency grid for the Wigner function evaluation (second output dimension).
Input	f_field (jnp.ndarray)	Frequency field $E(f)$.
Output	jnp.ndarray	The Wigner function array $W(t, f)$ with shape $(\text{len}(t), \text{len}(f))$.

Operation: The Wigner function is calculated using the definition via the temporal field $E(t)$: $$W(t, f) = \int d\tau' : E(t + \tau'/2) \cdot E^*(t - \tau'/2) \cdot e^{i 2\pi f \tau'}$$ This calculation is delegated to aux_functions.compute_wigner.

---

You've provided the content for two critical modules, pulses.py and aux_functions.py, which together handle the fundamental pulse mathematics for your JAX-based library.

Here is the complete documentation for both modules, structured with explanations and detailed input/output tables.

---

Module 1: pulses.py (Pulse Shape Definitions)

This module defines several common electric field pulse shapes in the frequency domain, normalized to represent their bandwidth based on the FWHM of the power profile ($|E(f)|^2$).

1. Curve Shape Definitions

gaussian_curve

@jit
def gaussian_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:

Defines a Gaussian pulse shape $E(f)$.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array [PHz].
Input	f_0 (float)	Pulse center frequency [PHz].
Input	gamma_0 (float)	Pulse bandwidth parameter related to the FWHM of $
Output	jnp.ndarray	Evaluated electric field in the frequency domain.

Operation: $$E(f) = \exp\left(-\frac{1}{2}\left(\frac{f-f_0}{\gamma_0}\right)^2\right)$$

lorentzian_curve

@jit
def lorentzian_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:

Defines a Lorentzian pulse shape $E(f)$.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array [PHz].
Input	f_0 (float)	Pulse center frequency [PHz].
Input	gamma_0 (float)	Pulse bandwidth parameter related to the FWHM of $
Output	jnp.ndarray	Evaluated electric field in the frequency domain.

Operation: $$E(f) = \frac{1}{1 + \left(\frac{f-f_0}{\gamma_0}\right)^2}$$

psquare_curve

@jit
def psquare_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float,
    smoothness: float = 0.005
) -> jnp.ndarray:

Defines a Smoothed Square (or Super-Gaussian) pulse shape $E(f)$ using a pair of sigmoid functions.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array [PHz].
Input	f_0 (float)	Pulse center frequency [PHz].
Input	gamma_0 (float)	Half-width parameter related to the FWHM of $
Input	smoothness (float)	Parameter controlling the sharpness of the edges (lower is sharper).
Output	jnp.ndarray	Evaluated electric field in the frequency domain.

Operation: The shape is defined by the product of a rising sigmoid (left edge) and a falling sigmoid (right edge) centered at $f_0 \pm \gamma_0$.

sech_curve

@jit
def sech_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:

Defines a Sech pulse shape $E(f)$.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array [PHz].
Input	f_0 (float)	Pulse center frequency [PHz].
Input	gamma_0 (float)	Pulse bandwidth parameter related to the FWHM of $
Output	jnp.ndarray	Evaluated electric field in the frequency domain.

Operation: $$E(f) = \text{sech}\left(\frac{f-f_0}{\gamma_0}\right) = \frac{1}{\cosh\left(\frac{f-f_0}{\gamma_0}\right)}$$

sech2_curve

@jit
def sech2_curve(
    f: jnp.ndarray,
    f_0: float,
    gamma_0: float
) -> jnp.ndarray:

Defines a Sech squared ($\text{Sech}^2$) pulse shape $E(f)$.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array [PHz].
Input	f_0 (float)	Pulse center frequency [PHz].
Input	gamma_0 (float)	Pulse bandwidth parameter related to the FWHM of $
Output	jnp.ndarray	Evaluated electric field in the frequency domain.

Operation: $$E(f) = \text{sech}^2\left(\frac{f-f_0}{\gamma_0}\right) = \frac{1}{\cosh^2\left(\frac{f-f_0}{\gamma_0}\right)}$$

2. Constants and Lookups

Name	Type	Description
gaussian	int (0)	Index for gaussian_curve.
lorentzian	int (1)	Index for lorentzian_curve.
square	int (2)	Index for psquare_curve.
sech	int (3)	Index for sech_curve.
sech2	int (4)	Index for sech2_curve.
curve_shapes	list	List mapping curve index to the JAX function.
fmwh_const	jnp.ndarray	Constants used to convert the input gamma_0 parameter into the exact FWHM of the power profile ($

---

---

Module 2: aux_functions.py (Mathematical Utilities)

This module provides utilities for numerical integration, vectorized operations, limit finding, and time-frequency analysis, essential for transforming and analyzing the electric field.

1. Fourier Transforms

scalar_inverse_fourier

@partial(jit) 
def scalar_inverse_fourier(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: float,
) -> float:

Calculates the inverse Fourier transform for a single point in time, $E(t)$. Used as the base function for the vectorized version.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array.
Input	f_field (jnp.ndarray)	Pulse in the frequency domain, $E(f)$.
Input	t (float)	Scalar point in time to evaluate the field.
Output	float	Evaluated pulse in the time domain, $E(t)$.

Operation: Performs a numerical Riemann sum approximation of the inverse Fourier integral: $$E(t) \approx \sum_f E(f) e^{-i 2\pi f t} \Delta f$$

vectorized_inverse_fourier

vectorized_inverse_fourier = vmap(
    scalar_inverse_fourier, 
    in_axes=(None, None, 0)
)

A vectorized version of scalar_inverse_fourier using JAX's vmap to efficiently compute $E(t)$ over an array of time points.

2. Instantaneous Frequency

scalar_instantaneous_frequency

@partial(jit)
def scalar_instantaneous_frequency(
    f: jnp.ndarray,
    f_field: jnp.ndarray,
    t: float,
) -> float:

Calculates the instantaneous frequency $\omega_{inst}(t)$ for a single time point.

Type	Name	Description
Input	f (jnp.ndarray)	Frequencies array.
Input	f_field (jnp.ndarray)	Pulse in the frequency domain, $E(f)$.
Input	t (float)	Scalar point in time to evaluate the field.
Output	float	Instantaneous frequency $\omega_{inst}(t)$ (in Hz).

Operation: Uses the definition based on the temporal field $E(t)$ and its time derivative $\partial_t E(t)$: $$\omega_{inst}(t) = -\frac{1}{2\pi} \text{Im}\left{\frac{\partial_t E(t)}{E(t)}\right}$$ The derivative $\partial_t E(t)$ is calculated in the frequency domain as the inverse Fourier transform of $-i 2\pi f E(f)$.

aux_instantaneous_frequency

aux_instantaneous_frequency = vmap(
    scalar_instantaneous_frequency,
    in_axes=(None, None, 0)
)

A vectorized version of scalar_instantaneous_frequency using JAX's vmap to compute the instantaneous frequency over an array of time points.

3. Numerical Limit Finding

set_limits

def set_limits(
    func: Callable,
    step: float,
    f_0: float,
    eps: float = 1e-5,
    max_steps = 10_000,
    negative: bool = False,
) -> float:

A generic numerical function to find the coordinate limit (time or frequency) where the normalized absolute value of a function (e.g., pulse field) drops below a specified tolerance ($\epsilon$).

Type	Name	Description
Input	func (Callable)	The function $F(x)$ to be evaluated (e.g., $
Input	step (float)	Size of the steps taken during the search.
Input	f_0 (float)	Center of the interval to start the search.
Input	eps (float)	Minimum normalized value to set the limit (e.g., $10^{-5}$ of the peak).
Input	max_steps (float)	Maximum number of iterations allowed.
Input	negative (bool)	If True, searches in the negative direction from $f_0$.
Output	float	The coordinate limit (distance from $f_0$) found, or the maximum distance if limit is not reached.

Operation: The function iteratively evaluates the ratio $|F(x)|/|F(f_0)|$ at increasing steps away from $f_0$. The process stops when the ratio is less than $\epsilon$ or $\text{max_steps}$ is reached.

4. Time-Frequency Analysis

compute_wigner

def compute_wigner(
    t_plot: jnp.ndarray,
    f_plot: jnp.ndarray,
    f_arr: jnp.ndarray,
    f_field: jnp.ndarray,
) -> jnp.ndarray:

Computes the Wigner function $W(t, f)$, a time-frequency phase space representation, using a memory-efficient lax.fori_loop structure vectorized over the 2D grid.

Type	Name	Description
Input	t_plot (jnp.ndarray)	Time grid for the Wigner function (first output dimension).
Input	f_plot (jnp.ndarray)	Frequency grid for the Wigner function (second output dimension).
Input	f_arr (jnp.ndarray)	Frequency array used for the inverse Fourier transforms (the grid of $E(f)$).
Input	f_field (jnp.ndarray)	The field in the frequency domain, $E(f)$.
Output	jnp.ndarray	Wigner function array $W(t, f)$ with shape $(\text{len}(t), \text{len}(f))$.

Operation: Calculates the integral definition of the Wigner function $W(t, f)$: $$W(t, f) = \int d\tau : E(t + \tau/2) \cdot E^*(t - \tau/2) \cdot e^{i 2\pi f \tau}$$ The computation is implemented using nested vmap and lax.fori_loop to perform the $\tau$ integration at each $(t, f)$ point efficiently.