array_tools/array_design/array_gen.py

import numpy as np
from typing import Tuple, Optional
from .base_array import *

def create_ula(
    elem_num: int,
    spacing: float,
    position_errors: Optional[np.ndarray] = None,
    gain_errors: Optional[np.ndarray] = None,
    phase_errors: Optional[np.ndarray] = None
) -> Array1D:
    """
    Create a Uniform Linear Array (ULA) along the x-axis with specified element count and spacing.

    Args:
        elem_num (int): Number of elements M (must be a positive integer).
        spacing (float): Distance between adjacent elements (in meters or wavelength units, must be > 0).
        position_errors (Optional[np.ndarray]): Optional perturbation offsets for each element, (M,), (2, M), or (3, M).
        gain_errors (Optional[np.ndarray]): Optional gain error factors (amplitude scale), shape (M,).
        phase_errors (Optional[np.ndarray]): Optional phase errors (in radians), shape (M,).

    Returns:
        Array1D: Initialized ULA object with 1D position array.

    Raises:
        ValueError: If elem_num is not a positive integer or spacing is not positive.

    Notes:
        - Element positions are on the x-axis only.
        - `positions` is a 1D NumPy array of shape (M,), i.e., scalar x-coordinates.
    """
    if not isinstance(elem_num, int) or elem_num <= 0:
        raise ValueError(f"elem_num must be a positive integer, got {elem_num}.")
    if not isinstance(spacing, (float, int)) or spacing <= 0:
        raise ValueError(f"spacing must be a positive number, got {spacing}.")

    # Uniform 1D positions: [0, spacing, ..., (M-1)*spacing]
    indices = np.arange(elem_num)
    positions = spacing * indices  # shape: (M,)

    return Array1D(
        positions=positions,
        type=ArrayType.ULA,
        spacing=spacing,
        position_errors=position_errors,
        gain_errors=gain_errors,
        phase_errors=phase_errors
    )

def create_ura(
    array_size: Tuple[int, int],
    spacing: float,
    position_errors: Optional[np.ndarray] = None,
    gain_errors: Optional[np.ndarray] = None,
    phase_errors: Optional[np.ndarray] = None
) -> Array2D:
    """
    Create a Uniform Rectangular Array (URA) in the xy-plane.

    The array consists of elements placed on a 2D rectangular grid, with uniform
    spacing in both x and y directions. The array lies entirely in the z=0 plane.

    Args:
        array_size (Tuple[int, int]): Number of elements along (cols, rows), i.e., (x-axis, y-axis).
                                      The total number of elements is M = cols * rows.
        spacing (float): Distance between adjacent elements (in meters or wavelength units).
        position_errors (Optional[np.ndarray]): Optional position perturbations of shape (2, M) or (3, M).
        gain_errors (Optional[np.ndarray]): Optional gain errors of shape (M,).
        phase_errors (Optional[np.ndarray]): Optional phase errors (in radians), shape (M,).

    Returns:
        Array2D: Initialized uniform rectangular array object.

    Raises:
        ValueError: If array dimensions are not positive integers or spacing is not positive.

    Notes:
        - Output positions are of shape (2, M), each column is [x; y] of an element.
        - Elements are ordered row-first: for each row, elements increase left to right (x).
        - The array is centered at the origin only if manually shifted after creation.
    """
    num_cols, num_rows = array_size
    if not isinstance(num_cols, int) or num_cols <= 0 or not isinstance(num_cols, int) or num_cols <= 0:
        raise ValueError(f"array_size must be two positive integers, got ({num_cols}, {num_cols}).")
    if not isinstance(spacing, (float, int)) or spacing <= 0:
        raise ValueError(f"spacing must be positive, got {spacing}.")

    total_elements = num_cols * num_rows

    # Generate 2D grid coordinates in xy-plane
    x_coords = spacing * np.arange(num_cols)  # x: fast-changing axis
    y_coords = spacing * np.arange(num_rows)  # y: slow-changing axis
    xv, yv = np.meshgrid(x_coords, y_coords)  # shape (cols, rows)

    # Flatten and stack into (2, M) position matrix
    positions = np.zeros((2, total_elements), dtype=float)
    positions[0, :] = xv.flatten()  # x-axis
    positions[1, :] = yv.flatten()  # y-axis

    return Array2D(
        positions=positions,
        type=ArrayType.URA,
        spacing=spacing,
        position_errors=position_errors,
        gain_errors=gain_errors,
        phase_errors=phase_errors
    )

def create_uca(
    elem_num: int,
    radius: float,
    position_errors: Optional[np.ndarray] = None,
    gain_errors: Optional[np.ndarray] = None,
    phase_errors: Optional[np.ndarray] = None
) -> Array2D:
    """
    Create a Uniform Circular Array (UCA) centered at the origin in the xy-plane.

    The array elements are evenly spaced along the circumference of a circle,
    with the first element at angle 0 (x-axis), and proceeding counterclockwise.

    Args:
        elem_num (int): Number of elements M (must be a positive integer).
        radius (float): Radius of the circular array (in meters or wavelength units, must be > 0).
        position_errors (Optional[np.ndarray]): Optional position perturbations, shape (2, M) or (3, M).
        gain_errors (Optional[np.ndarray]): Optional gain error factors (amplitude scale), shape (M,).
        phase_errors (Optional[np.ndarray]): Optional phase errors (in radians), shape (M,).

    Returns:
        Array2D: Initialized array object with circular geometry.

    Raises:
        ValueError:
            - If elem_num is not a positive integer.
            - If radius is not a positive number.

    Notes:
        - The output positions are of shape (2, M), representing [x; y] coordinates.
        - The nominal element spacing (along the arc) is returned as the 'spacing' field.
        - The array lies entirely in the xy-plane (z = 0 is implied).
    """
    if not isinstance(elem_num, int) or elem_num <= 0:
        raise ValueError(f"elem_num ({elem_num}) must be a positive integer.")
    if not isinstance(radius, (float, int)) or radius <= 0:
        raise ValueError(f"radius ({radius}) must be a positive number.")

    # Compute angular positions (uniform spacing)
    theta = np.linspace(0, 2 * np.pi, elem_num, endpoint=False)

    # Compute Cartesian coordinates in 2D
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)

    # Store as 2×N position matrix: positions[0, :] = x, positions[1, :] = y
    positions = np.zeros((2, elem_num), dtype=float)
    positions[0, :] = x
    positions[1, :] = y

    # Compute approximate arc-length spacing between adjacent elements
    step = 2 * np.pi / elem_num
    spacing = 2 * radius * np.sin(step / 2)  # chord length ≈ arc length for small step

    return Array2D(
        positions=positions,
        type=ArrayType.UCA,
        spacing=spacing,
        position_errors=position_errors,
        gain_errors=gain_errors,
        phase_errors=phase_errors
    )

def create_cyl_array_with_uca(
    elem_per_layer: int,
    radius: float,
    num_layers: int,
    layer_spacing: float,
    gain_errors: Optional[np.ndarray] = None,
    phase_errors: Optional[np.ndarray] = None,
    position_errors: Optional[np.ndarray] = None
) -> Array3D:
    """
    Create a cylindrical array composed of stacked UCAs along the z-axis.

    Args:
        elem_per_layer (int): Number of elements per circular UCA layer (must be > 0).
        radius (float): Radius of each UCA layer (must be > 0).
        num_layers (int): Number of stacked layers along z-axis (must be > 0).
        layer_spacing (float): Distance between adjacent layers along z-axis (must be > 0).
        gain_errors (Optional[np.ndarray]): Optional gain errors, shape (M,).
        phase_errors (Optional[np.ndarray]): Optional phase errors, shape (M,).
        position_errors (Optional[np.ndarray]): Optional position errors, shape (3, M).

    Returns:
        Array3D: Array object with cylindrical geometry.

    Raises:
        ValueError: If any input is invalid (non-positive integer or float).
    """
    if elem_per_layer <= 0 or num_layers <= 0:
        raise ValueError("elem_per_layer and num_layers must be positive integers.")
    if radius <= 0:
        raise ValueError("radius must be a positive number.")
    if layer_spacing <= 0:
        raise ValueError("layer_spacing must be a positive number.")

    total_elements = elem_per_layer * num_layers
    positions = np.zeros((3, total_elements), dtype=float)

    theta = np.linspace(0, 2 * np.pi, elem_per_layer, endpoint=False)
    x = radius * np.cos(theta)
    y = radius * np.sin(theta)

    for i in range(num_layers):
        z = layer_spacing * i
        idx_start = i * elem_per_layer
        idx_end = (i + 1) * elem_per_layer
        positions[0, idx_start:idx_end] = x
        positions[1, idx_start:idx_end] = y
        positions[2, idx_start:idx_end] = z

    return Array3D(
        positions=positions,
        type=ArrayType.CYLINDRICAL,
        spacing=layer_spacing,
        position_errors=position_errors,
        gain_errors=gain_errors,
        phase_errors=phase_errors
    )