1#!/usr/bin/env python3
2"""
3Z-Transform Analysis
4====================
5
6This script demonstrates the Z-transform and its applications in discrete-time
7signal processing, including:
8
9- Common Z-transform pairs and their verification
10- Pole-zero diagrams with unit circle and stability regions
11- System stability analysis from pole locations
12- Inverse Z-transform via partial fraction expansion
13- 2nd-order IIR system analysis (poles, zeros, frequency response)
14
15Key Concepts:
16 X(z) = sum_{n=-inf}^{inf} x[n] * z^{-n}
17 - Poles inside the unit circle โ stable causal system
18 - Poles outside the unit circle โ unstable causal system
19 - Poles on the unit circle โ marginally stable
20
21Author: Educational example for Signal Processing
22License: MIT
23"""
24
25import numpy as np
26from scipy import signal
27import matplotlib.pyplot as plt
28
29# ============================================================================
30# Z-TRANSFORM UTILITIES
31# ============================================================================
32
33def compute_z_transform_samples(x, z_values):
34 """
35 Numerically evaluate X(z) = sum_{n=0}^{N-1} x[n] * z^{-n}
36 for a finite-length sequence x[n].
37
38 Args:
39 x (ndarray): Input sequence (assumed causal, starting at n=0)
40 z_values (ndarray): Complex z values at which to evaluate X(z)
41
42 Returns:
43 ndarray: X(z) evaluated at each z value
44 """
45 N = len(x)
46 n = np.arange(N)
47 # X(z) = sum x[n] * z^{-n} โ broadcast over z_values
48 # Shape: (len(z_values), N), then sum along axis=1
49 z_inv_n = z_values[:, np.newaxis] ** (-n[np.newaxis, :])
50 return (x[np.newaxis, :] * z_inv_n).sum(axis=1)
51
52
53def partial_fraction_inverse(b, a, n_samples=20):
54 """
55 Compute inverse Z-transform via partial fraction expansion.
56
57 Uses scipy.signal.residuez to decompose H(z) = B(z)/A(z) into:
58 H(z) = sum_k r_k / (1 - p_k * z^{-1}) + direct terms
59
60 The inverse Z-transform of r_k / (1 - p_k * z^{-1}) is r_k * p_k^n * u[n].
61
62 Args:
63 b (array-like): Numerator polynomial coefficients (descending z powers)
64 a (array-like): Denominator polynomial coefficients
65 n_samples (int): Number of samples for h[n]
66
67 Returns:
68 tuple: (h, residues, poles, direct_terms)
69 """
70 # residuez works with z^{-1} polynomials (signal convention)
71 r, p, k = signal.residuez(b, a)
72
73 # Reconstruct h[n] = sum_i r[i] * p[i]^n for n >= 0
74 # The imaginary parts of complex-conjugate pairs cancel exactly;
75 # taking np.real() discards the numerical rounding residual.
76 n = np.arange(n_samples)
77 h = np.zeros(n_samples, dtype=complex)
78 for ri, pi in zip(r, p):
79 h += ri * pi ** n
80 h = np.real(h) # imaginary part is ~machine-epsilon for real-valued systems
81
82 return h, r, p, k
83
84
85# ============================================================================
86# MAIN DEMONSTRATIONS
87# ============================================================================
88
89def demo_common_z_transforms():
90 """Verify common Z-transform pairs numerically."""
91 print("=" * 65)
92 print("1. COMMON Z-TRANSFORM PAIRS โ NUMERICAL VERIFICATION")
93 print("=" * 65)
94
95 # Evaluate on a circle of radius r = 1.2 (inside ROC for stable signals)
96 theta = np.linspace(0, 2 * np.pi, 500, endpoint=False)
97 r = 1.2
98 z = r * np.exp(1j * theta)
99
100 # --- Pair 1: Unit step u[n] โ z / (z - 1), |z| > 1
101 N = 200 # approximate infinite sum with long finite sequence
102 x_step = np.ones(N)
103 X_numerical = compute_z_transform_samples(x_step, z)
104 X_analytical = z / (z - 1)
105 err_step = np.max(np.abs(X_numerical - X_analytical)) / np.max(np.abs(X_analytical))
106 print(f"\nUnit step u[n] โ z/(z-1)")
107 print(f" Max relative error (N={N} terms): {err_step:.4e}")
108
109 # --- Pair 2: Exponential a^n * u[n] โ z / (z - a), |z| > |a|
110 a = 0.8
111 x_exp = a ** np.arange(N)
112 X_numerical = compute_z_transform_samples(x_exp, z)
113 X_analytical = z / (z - a)
114 err_exp = np.max(np.abs(X_numerical - X_analytical)) / np.max(np.abs(X_analytical))
115 print(f"\nExponential a^nยทu[n] (a={a}) โ z/(z-a)")
116 print(f" Max relative error (N={N} terms): {err_exp:.4e}")
117
118 # --- Pair 3: Unit impulse ฮด[n] โ 1 (all z)
119 x_imp = np.zeros(50); x_imp[0] = 1.0
120 X_numerical = compute_z_transform_samples(x_imp, z)
121 X_analytical = np.ones_like(z)
122 err_imp = np.max(np.abs(X_numerical - X_analytical))
123 print(f"\nUnit impulse ฮด[n] โ 1")
124 print(f" Max absolute error: {err_imp:.4e}")
125
126 # --- Pair 4: Delayed impulse ฮด[n-k] โ z^{-k}
127 k_delay = 3
128 x_del = np.zeros(50); x_del[k_delay] = 1.0
129 X_numerical = compute_z_transform_samples(x_del, z)
130 X_analytical = z ** (-k_delay)
131 err_del = np.max(np.abs(X_numerical - X_analytical))
132 print(f"\nDelayed impulse ฮด[n-{k_delay}] โ z^{{-{k_delay}}}")
133 print(f" Max absolute error: {err_del:.4e}")
134
135
136def demo_pole_zero_and_stability():
137 """Plot pole-zero diagrams and assess stability."""
138 print("\n" + "=" * 65)
139 print("2. POLE-ZERO DIAGRAMS AND STABILITY ANALYSIS")
140 print("=" * 65)
141
142 # System 1: Stable โ poles inside unit circle
143 b1 = [1, -0.5] # zeros of H1(z)
144 a1 = [1, -0.9, 0.2] # poles from denominator
145 z1, p1, k1 = signal.tf2zpk(b1, a1)
146 stable1 = np.all(np.abs(p1) < 1.0)
147 print(f"\nSystem 1 H(z) = (z - 0.5) / (z^2 - 0.9z + 0.2)")
148 print(f" Zeros : {z1}")
149 print(f" Poles : {p1}")
150 print(f" |Poles|: {np.abs(p1)}")
151 print(f" Stable? {stable1} (all poles inside unit circle)")
152
153 # System 2: Unstable โ one pole outside unit circle
154 b2 = [1, 0]
155 a2 = [1, -1.5, 0.56] # roots at 0.8 and 0.7 โ stable
156 # Let's make one pole outside: a2 roots at 1.2 and 0.5
157 a2_unstable = np.poly([1.2, 0.5])
158 z2, p2, k2 = signal.tf2zpk(b2, a2_unstable)
159 stable2 = np.all(np.abs(p2) < 1.0)
160 print(f"\nSystem 2 (one pole outside unit circle)")
161 print(f" Poles : {p2}")
162 print(f" |Poles|: {np.abs(p2)}")
163 print(f" Stable? {stable2}")
164
165 # System 3: Marginally stable โ complex poles ON unit circle
166 omega0 = np.pi / 4 # oscillation frequency
167 p3 = np.array([np.exp(1j * omega0), np.exp(-1j * omega0)])
168 b3, a3 = signal.zpk2tf([], p3, 1.0)
169 z3, p3_check, _ = signal.tf2zpk(b3, a3)
170 print(f"\nSystem 3 (complex poles on unit circle, ฯโ=ฯ/4)")
171 print(f" Poles : {np.round(p3_check, 4)}")
172 print(f" |Poles|: {np.abs(p3_check)}")
173 print(f" Marginally stable (oscillator, never decays or grows)")
174
175 return (b1, a1, z1, p1), (b2, a2_unstable, z2, p2), (b3, a3, z3, p3_check)
176
177
178def demo_2nd_order_iir():
179 """Full analysis of a 2nd-order IIR system via Z-transform."""
180 print("\n" + "=" * 65)
181 print("3. 2nd-ORDER IIR SYSTEM ANALYSIS")
182 print("=" * 65)
183
184 # H(z) = (1 - 0.5 z^{-1}) / (1 - 0.6 z^{-1} + 0.5 z^{-2})
185 # Written in positive powers: H(z) = z(z - 0.5) / (z^2 - 0.6z + 0.5)
186 b = [1, -0.5, 0] # numerator coefficients (z^2, z^1, z^0)
187 a = [1, -0.6, 0.5] # denominator coefficients
188
189 zeros, poles, gain = signal.tf2zpk(b, a)
190 print(f"\nH(z) = (zยฒ - 0.5z) / (zยฒ - 0.6z + 0.5)")
191 print(f" Zeros : {np.round(zeros, 4)}")
192 print(f" Poles : {np.round(poles, 4)}")
193 print(f" |Poles|: {np.round(np.abs(poles), 4)} โ stable: {np.all(np.abs(poles) < 1)}")
194
195 # Inverse Z-transform via partial fractions
196 h_pf, residues, pf_poles, direct = partial_fraction_inverse(b, a, n_samples=30)
197 print(f"\nPartial Fraction Expansion:")
198 for i, (ri, pi) in enumerate(zip(residues, pf_poles)):
199 print(f" r_{i} = {ri:.4f}, p_{i} = {pi:.4f} โ term: {ri:.4f} * ({pi:.4f})^n")
200
201 # Compare with direct impulse response from lfilter
202 impulse = np.zeros(30); impulse[0] = 1.0
203 h_direct = signal.lfilter(b, a, impulse)
204 err = np.max(np.abs(h_pf - h_direct))
205 print(f"\nImpulse response comparison (PF vs lfilter): max error = {err:.2e}")
206
207 # Frequency response
208 w, H = signal.freqz(b, a, worN=512)
209 print(f"\nFrequency response (first 5 points):")
210 for i in range(5):
211 print(f" ฯ = {w[i]:.3f} rad/sample |H| = {np.abs(H[i]):.4f}")
212
213 return b, a, zeros, poles, w, H, h_direct
214
215
216# ============================================================================
217# VISUALIZATION
218# ============================================================================
219
220def plot_pole_zero(ax, zeros, poles, title):
221 """Draw a pole-zero diagram on the given axes."""
222 # Unit circle
223 theta = np.linspace(0, 2 * np.pi, 300)
224 ax.plot(np.cos(theta), np.sin(theta), 'k--', linewidth=0.8, alpha=0.5, label='Unit circle')
225 ax.axhline(0, color='k', linewidth=0.5)
226 ax.axvline(0, color='k', linewidth=0.5)
227
228 # Zeros (circles) and poles (crosses)
229 if len(zeros) > 0:
230 ax.scatter(np.real(zeros), np.imag(zeros), marker='o', s=80,
231 facecolors='none', edgecolors='blue', linewidths=2, zorder=5, label='Zeros')
232 if len(poles) > 0:
233 ax.scatter(np.real(poles), np.imag(poles), marker='x', s=80,
234 color='red', linewidths=2, zorder=5, label='Poles')
235
236 ax.set_xlabel('Real Part')
237 ax.set_ylabel('Imaginary Part')
238 ax.set_title(title)
239 ax.legend(fontsize=8)
240 ax.set_aspect('equal')
241 ax.grid(True, alpha=0.3)
242 lim = 1.6
243 ax.set_xlim(-lim, lim)
244 ax.set_ylim(-lim, lim)
245
246
247if __name__ == "__main__":
248 print("Z-TRANSFORM ANALYSIS")
249 print("=" * 65)
250
251 # --- Numerical demonstrations ---
252 demo_common_z_transforms()
253 sys1, sys2, sys3 = demo_pole_zero_and_stability()
254 b, a, zeros, poles, w, H, h_direct = demo_2nd_order_iir()
255
256 # --- Visualizations ---
257 fig, axes = plt.subplots(2, 3, figsize=(15, 9))
258 fig.suptitle('Z-Transform Analysis', fontsize=14, fontweight='bold')
259
260 # Row 0: Pole-zero diagrams for the three systems
261 plot_pole_zero(axes[0, 0], sys1[2], sys1[3], 'Stable System\n(poles inside unit circle)')
262 plot_pole_zero(axes[0, 1], sys2[2], sys2[3], 'Unstable System\n(pole outside unit circle)')
263 plot_pole_zero(axes[0, 2], sys3[2], sys3[3], 'Marginally Stable\n(poles on unit circle)')
264
265 # Row 1, col 0: Pole-zero of 2nd-order IIR system
266 plot_pole_zero(axes[1, 0], zeros, poles, '2nd-Order IIR\nPole-Zero Diagram')
267
268 # Row 1, col 1: Impulse response (inverse Z-transform)
269 n_plot = np.arange(30)
270 axes[1, 1].stem(n_plot, h_direct, basefmt='k-', linefmt='C0-', markerfmt='C0o')
271 axes[1, 1].set_xlabel('Sample n')
272 axes[1, 1].set_ylabel('h[n]')
273 axes[1, 1].set_title('Impulse Response h[n]\n(Inverse Z-Transform via PFE)')
274 axes[1, 1].grid(True, alpha=0.3)
275
276 # Row 1, col 2: Magnitude and phase frequency response
277 axes[1, 2].plot(w / np.pi, 20 * np.log10(np.maximum(np.abs(H), 1e-10)),
278 'b-', linewidth=1.8)
279 axes[1, 2].set_xlabel('Normalized Frequency (รฯ rad/sample)')
280 axes[1, 2].set_ylabel('Magnitude (dB)')
281 axes[1, 2].set_title('2nd-Order IIR Frequency Response\n(from Z-Transform)')
282 axes[1, 2].grid(True, alpha=0.3)
283
284 plt.tight_layout()
285 out_path = '/opt/projects/01_Personal/03_Study/examples/Signal_Processing/07_z_transform.png'
286 plt.savefig(out_path, dpi=150)
287 print(f"\nSaved visualization: {out_path}")
288 plt.close()
289
290 print("\n" + "=" * 65)
291 print("Key Takeaways:")
292 print(" - ROC determines whether the Z-transform exists")
293 print(" - Poles inside unit circle โ stable causal system")
294 print(" - Partial fraction expansion gives the inverse Z-transform")
295 print(" - Z-transform on the unit circle equals the DTFT")
296 print("=" * 65)