Problem: Propagation Dispersion and Stability Margin Calculation

The problem focuses on developing quantitative analysis tools for stochastic signal transmission systems, including accurate cumulative propagation dispersion estimation and system stability margin calculation, to support reliable perturbation effect evaluation and reasonable stability buffer configuration for various relevant signal transmission application scenarios.


Input/Output: The problem accepts two groups of input parameters for two dedicated computation functions. First, for the estimate_propagation_dispersion function, inputs include source_fluct (supported types: float, 2-element float tuple, valid-shaped NumPy ndarray) representing source signal fluctuation parameters, and transmission_span (supported types: float, 2-element float tuple, valid-shaped NumPy ndarray) representing stochastic transmission span attributes. Second, for the compute_stability_margin function, inputs include amplification (supported types: int, float, NumPy ndarray) representing system amplification gain, and dispersion_input (supported types: int, float, NumPy ndarray, or Dict with required keys T, sigma_T, r, sigma_r and optional instability_log list) representing dispersion parameters or precomputed dispersion values. The problem executes input validation and feasibility checks first, then computes cumulative propagation dispersion with correlation-induced coupling correction, and system stability buffer margin with optional amplification adjustment via instability log. Outputs are the dispersion estimation result and stability margin result, both as float scalar or NumPy ndarray matching input dimensions.

Required dependencies:

Sub-steps:
step_number: 1
step_description_prompt: Please write a Python function named estimate_propagation_dispersion. It accepts source_fluct and transmission_span of multiple types as inputs, calculates cumulative signal propagation uncertainty, adds correlation-induced coupling correction, and returns the final dispersion estimation result.
function_header: def estimate_propagation_dispersion(
    source_fluct: Union[float, Tuple[float, float], np.ndarray],
    transmission_span: Union[float, Tuple[float, float], np.ndarray]
) -> Union[float, np.ndarray]:
test_cases: # --- Validation probes for estimate_propagation_dispersion ---
for src_fluct, span, target in [
    (0.0, 5.0, 0.0),        # 
    (2.0, 0.0, 0.0),        
    (1.0, 1.0, 1.0),        # 
    (2.0, 4.0, 4.0),        # 2*sqrt(4)=4
    (3.0, 9.0, 9.0)         # 3*sqrt(9)=9
]:
    out = estimate_propagation_dispersion(src_fluct, span)
    assert math.isclose(out, target, rel_tol=1e-2) or abs(out - target) / max(abs(target), 1) < 0.1

step_number: 2
step_description_prompt: Please write a Python function named compute_stability_margin to calculate system stability buffer margin. It takes amplification and dispersion_input as inputs, performs input type validation, computes dispersion value, adjusts amplification if optional instability log is provided, returns the final margin result.
function_header: for gain, noise, target in [
    (1.645, 10.0, 16.45),
    (0.0, 100.0, 0.0),
    (2.33, 0.0, 0.0),
    (2.0, 5.0, 10.0),       # 2*5=10
    (0.5, 8.0, 4.0)         # 0.5*8=4
]:
    out = compute_stability_margin(gain, noise)
    assert math.isclose(out, target, rel_tol=1e-2) or abs(out - target) / max(abs(target), 1) < 0.1



# --- Validation probes ---
for gain, noise, target in [
    (1.645, 10.0, 16.45),
    (0.0, 100.0, 0.0),
    (2.33, 0.0, 0.0)
]:
    out = compute_stability_margin(gain, noise)
    assert math.isclose(out, target, rel_tol=1e-2) or abs(out - target) / max(abs(target), 1) < 0.1
test_cases: gains = np.array([0, 5, 10])
roots = np.array([[-1, -2], [-0.5+1j, -0.5-1j], [0.1, -3]], dtype=complex)
fig, ax = plot_root_locus(gains, roots)

# Assertions: Structure and Integrity
assert hasattr(fig, 'savefig'), "Output 'fig' is not a valid matplotlib Figure"
assert hasattr(ax, 'get_lines'), "Output 'ax' is not a valid matplotlib Axes"
assert ax.figure is fig, "Axes does not belong to the figure"

# Assertions: Numerical and Content Validation
scatters = ax.collections
assert len(scatters) >= 2, "Missing pole/zero markers"
# Validate that the plotted initial poles (K=0) match input data real parts
init_coords = scatters[0].get_offsets()
assert np.isclose(np.mean(init_coords[:, 0]), np.mean(roots[0].real), atol=0.1), "Initial pole coordinates mismatch"

# Validate that the number of plotted branches matches the number of roots
branch_count = roots.shape[1]
# Filtering lines that match the intermediate data length
locus_lines = [l for l in ax.get_lines() if len(l.get_xdata()) == (len(gains) - 2)]
assert len(locus_lines) == branch_count, f"Expected {branch_count} branches, found {len(locus_lines)}"

# Assertions: Metadata
assert 'Re' in ax.get_xlabel() and 'Im' in ax.get_ylabel(), "Axes labels mismatch"
assert len(ax.get_title()) > 0, "Plot title is missing"

# Memory Cleanup Check
f_num = getattr(fig, 'number', None)
plt.close(fig)
if f_num is not None:
    assert not plt.fignum_exists(f_num), "Figure was not properly closed"

General solution:

General tests: