josh/array_tools
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: initial commit

Browse Source
This commit is contained in:
josh
2025-07-01 16:09:27 +08:00
commit bf8d7fd556
5 changed files with 682 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+195
View File
@@ -0,0 +1,195 @@
# Byte-compiled / optimized / DLL files
__pycache__/
*.py[cod]
*$py.class
# C extensions
*.so
# Distribution / packaging
.Python
build/
develop-eggs/
dist/
downloads/
eggs/
.eggs/
lib/
lib64/
parts/
sdist/
var/
wheels/
share/python-wheels/
*.egg-info/
.installed.cfg
*.egg
MANIFEST
# PyInstaller
# Usually these files are written by a python script from a template
# before PyInstaller builds the exe, so as to inject date/other infos into it.
*.manifest
*.spec
# Installer logs
pip-log.txt
pip-delete-this-directory.txt
# Unit test / coverage reports
htmlcov/
.tox/
.nox/
.coverage
.coverage.*
.cache
nosetests.xml
coverage.xml
*.cover
*.py,cover
.hypothesis/
.pytest_cache/
cover/
# Translations
*.mo
*.pot
# Django stuff:
*.log
local_settings.py
db.sqlite3
db.sqlite3-journal
# Flask stuff:
instance/
.webassets-cache
# Scrapy stuff:
.scrapy
# Sphinx documentation
docs/_build/
# PyBuilder
.pybuilder/
target/
# Jupyter Notebook
.ipynb_checkpoints
# IPython
profile_default/
ipython_config.py
# pyenv
# For a library or package, you might want to ignore these files since the code is
# intended to run in multiple environments; otherwise, check them in:
# .python-version
# pipenv
# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
# However, in case of collaboration, if having platform-specific dependencies or dependencies
# having no cross-platform support, pipenv may install dependencies that don't work, or not
# install all needed dependencies.
#Pipfile.lock
# UV
# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
#uv.lock
# poetry
# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
# This is especially recommended for binary packages to ensure reproducibility, and is more
# commonly ignored for libraries.
# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
#poetry.lock
#poetry.toml
# pdm
# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
#pdm.lock
# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
# in version control.
# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
.pdm.toml
.pdm-python
.pdm-build/
# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
__pypackages__/
# Celery stuff
celerybeat-schedule
celerybeat.pid
# SageMath parsed files
*.sage.py
# Environments
.env
.venv
env/
venv/
ENV/
env.bak/
venv.bak/
# Spyder project settings
.spyderproject
.spyproject
# Rope project settings
.ropeproject
# mkdocs documentation
/site
# mypy
.mypy_cache/
.dmypy.json
dmypy.json
# Pyre type checker
.pyre/
# pytype static type analyzer
.pytype/
# Cython debug symbols
cython_debug/
# PyCharm
# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
# and can be added to the global gitignore or merged into this file. For a more nuclear
# option (not recommended) you can uncomment the following to ignore the entire idea folder.
.idea/
# Abstra
# Abstra is an AI-powered process automation framework.
# Ignore directories containing user credentials, local state, and settings.
# Learn more at https://abstra.io/docs
.abstra/
# Visual Studio Code
# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
# and can be added to the global gitignore or merged into this file. However, if you prefer,
# you could uncomment the following to ignore the entire vscode folder
.vscode/
# Ruff stuff:
.ruff_cache/
# PyPI configuration file
.pypirc
# Cursor
# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
# refer to https://docs.cursor.com/context/ignore-files
.cursorignore
.cursorindexingignore
array_design/__init__.py
View File
array_design/base_array.py
+472
View File
@@ -0,0 +1,472 @@
from abc import ABC, abstractmethod
from enum import Enum
import numpy as np
from typing import Tuple, Optional
class ArrayType(Enum):
"""Enumeration of supported antenna array types."""
ULA = 'ula'
URA = 'ura'
UCA = 'uca'
NESTED = 'nested'
COPRIME = 'coprime'
MRA = 'mra'
CUSTOM = 'custom'
class BaseArray(ABC):
"""Abstract base class for antenna arrays.
Attributes:
positions (np.ndarray): Physical coordinates of each element, shape (dim, M).
- dim: Spatial dimension of the array (1 for linear, 2 for planar, etc.)
- M: Number of antenna elements
type (ArrayType): Array geometry type (e.g., ULA, URA, UCA, etc.).
spacing (Optional[float]): Inter-element spacing, if uniformly defined.
dim (int): Spatial dimension, inferred from positions.shape[0].
count (int): Total number of elements, inferred from positions.shape[1].
indices (Optional[np.ndarray]): Structured grid indices, shape (dim, M), used in structured arrays like URA.
position_errors (np.ndarray): Perturbations in element positions, same shape as `positions`.
gain_errors (np.ndarray): Per-element gain variation factors (linear scale), same shape as `positions`.
phase_errors (np.ndarray): Per-element phase offsets (in radians), same shape as `positions`.
"""
def __init__(
self,
positions: np.ndarray,
type: ArrayType,
spacing: Optional[float] = None,
indices: Optional[np.ndarray] = None,
position_errors: Optional[np.ndarray] = None,
gain_errors: Optional[np.ndarray] = None,
phase_errors: Optional[np.ndarray] = None
):
"""Initialize a base antenna array.
Args:
positions (np.ndarray): Element positions, shape (dim, M), where:
- dim is the spatial dimension (1, 2, or 3),
- M is the number of elements.
type (ArrayType): Type of the array structure (e.g., ArrayType.ULA, ArrayType.URA).
spacing (Optional[float]): Inter-element spacing, relevant only for structured arrays.
indices (Optional[np.ndarray]): Grid-based indices, same shape as `positions`.
position_errors (Optional[np.ndarray]): Position perturbations, same shape as `positions`.
gain_errors (Optional[np.ndarray]): Gain perturbations (linear), same shape as `positions`.
phase_errors (Optional[np.ndarray]): Phase perturbations (in radians), same shape as `positions`.
Raises:
ValueError: If input arrays have mismatched shapes or invalid dimensions.
"""
# position-related
positions = np.array(positions)
if positions.ndim != 2:
raise ValueError(f"`positions` must be a 2D array, got ndim={positions.ndim}.")
if positions.shape[0] not in (1, 2, 3):
raise ValueError(f"`positions` must be in 1D, 2D, or 3D space. Got dim={positions.shape[0]}.")
self.positions = positions
self.dim = positions.shape[0]
self.count = positions.shape[1]
# index
if indices is not None:
indices = np.array(indices)
if indices.shape != positions.shape:
raise ValueError(f"Shape mismatch: indices {indices.shape} vs positions {positions.shape}")
self.indices = indices
# errors and perturbations
if position_errors is not None:
position_errors = np.array(position_errors)
if position_errors.shape != positions.shape:
raise ValueError(f"Shape mismatch: position_errors {position_errors.shape} vs positions {positions.shape}")
self.position_errors = position_errors
else:
self.position_errors = np.zeros_like(positions)
if gain_errors is not None:
gain_errors = np.array(gain_errors)
if gain_errors.shape != positions.shape:
raise ValueError(f"Shape mismatch: gain_errors {gain_errors.shape} vs positions {positions.shape}")
self.gain_errors = gain_errors
else:
self.gain_errors = np.zeros_like(positions)
if phase_errors is not None:
phase_errors = np.array(phase_errors)
if phase_errors.shape != positions.shape:
raise ValueError(f"Shape mismatch: phase_errors {phase_errors.shape} vs positions {positions.shape}")
self.phase_errors = phase_errors
else:
self.phase_errors = np.zeros_like(positions)
# other attributes
self.spacing = spacing
self.type = type
@abstractmethod
def steering_matrix(
self,
wavelength: float,
directions: np.ndarray
) -> np.ndarray:
"""Computes the array steering matrix for specified direction(s).
Args:
wavelength (float): Wavelength of the incoming signal (same unit as element positions).
directions (np.ndarray): Incident direction(s) in radians. Shape should be:
- (dim,) for a single direction
- (dim, D) for D directions
Returns:
np.ndarray: Complex-valued steering matrix of shape (M, D), where M is the number of elements
and D is the number of directions. For a single direction, shape is (M,).
Raises:
ValueError: If `wavelength` is non-positive.
ValueError: If `directions` shape is incompatible with array dimension.
"""
pass
@abstractmethod
def rotate(
self,
target_normal: np.ndarray,
current_normal: Optional[np.ndarray] = None,
center: Optional[np.ndarray] = None
) -> "BaseArray":
"""Rotate the array such that its current normal vector aligns with a target direction.
This method performs a rigid-body rotation around a given center (default: centroid):
1. Translates the array such that the rotation center lies at the origin.
2. Applies the rotation that maps `current_normal` to `target_normal`.
3. Translates the array back to its centroid or specified center.
Args:
target_normal (np.ndarray): Target normal vector, shape (dim,). Must be non-zero.
current_normal (Optional[np.ndarray]): Current normal vector, (shape (dim,).
If None, defaults to the vector from centroid to origin.
center (Optional[np.ndarray]): Optional rotation center, shape (dim,).
If None, uses the centroid of the array.
Returns:
BaseArray: A new array instance with rotated element positions (and other attributes if applicable).
Raises:
ValueError: If `target_normal` or `current_normal` has incorrect shape or is near-zero in norm.
NotImplementedError: If the method is not implemented in a concrete subclass.
"""
pass
def translate(self, offset: np.ndarray):
"""Translate the array to a new center by applying the given offset.
Args:
offset (np.ndarray): Translation vector of shape (dim,) or (1, dim).
Raises:
ValueError: If the dimension of the offset does not match the array's spatial dimension.
"""
if offset.ndim == 1 and offset.shape[0] == self.dim:
offset = np.reshape(offset, (-1, 1))
elif offset.ndim == 2:
if offset.shape == (1, self.dim):
offset = np.reshape(offset, (self.dim, 1))
elif offset.shape == (self.dim, 1):
pass # already correct shape
else:
raise ValueError(
f"Offset must be convertible to shape ({self.dim}, 1), "
f"but got shape {offset.shape}."
)
else:
raise ValueError(
f"Offset must be a 1D or 2D array convertible to shape ({self.dim}, 1), "
f"but got shape {offset.shape}."
)
self.positions += offset
def center(self):
"""Center the array at the origin."""
centroid = np.mean(self.positions, axis=1)
self.translate(-centroid)
class Array1D(BaseArray):
"""
One-dimensional linear array.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if self.dim != 1:
raise ValueError(f"Array1D requires 1D array. Got dimension = {self.dim}.")
def steering_matrix(
self,
wavelength: float,
directions: np.ndarray,
compute_derivative: bool = False
) -> Tuple[np.ndarray, Optional[np.ndarray]]:
"""Compute the steering matrix of 1D array.
Args:
wavelength (float): Signal wavelength.
directions (np.ndarray): Input directions, must be in one of the following formats:
- shape (D,) or (1, D): D broadside angles (in radians), valid for 1D array
- shape (2, D): D pairs of azimuth and elevation angles (in radians), used for projection to broadside
The final broadside angle is computed as asin(sin(az) * cos(el)).
compute_derivative (bool): Whether to compute the derivative of the steering matrix w.r.t. each DOA (only valid for 1D DOAs).
Returns:
A (np.ndarray): Steering matrix of shape (M, D)
DA (np.ndarray or None): Derivative matrix (optional, only valid for 1D DOAs)
Raises:
ValueError: If dimension of directions is not 1 or 2.
"""
directions = np.array(directions)
if directions.ndim == 1:
directions = np.reshape(directions, (1, -1))
elif directions.ndim == 2:
if directions.shape[0] != 2:
raise ValueError(
f"If directions is 2D, its shape must be (2, D), "
f"but got shape {directions.shape}."
)
az, el = directions[0, :], directions[1, :]
directions = np.arcsin(np.sin(az) * np.cos(el))
else:
raise ValueError(
f"Invalid shape for directions ({directions.ndim}). "
"Expected shape (D,) or (1, D) for 1D broadside angles or shape (2, D) for azimuth/elevation."
)
DA = None
positions = self.positions
k = 2 * np.pi / wavelength
phase = k * (np.outer(positions, np.sin(directions)))
A = np.exp(1j * phase)
if compute_derivative:
if directions.ndim != 1:
raise ValueError("compute_derivative is only valid for 1D DOAs (broadside angles).")
DA = 1j * 2 * np.pi / wavelength * np.outer(positions, np.cos(directions)) * A
# apply gain/phase errors
A *= np.exp(1j * self.phase_errors.flatten())
A *= self.gain_errors.flatten()
if DA is not None:
A_phase = np.exp(1j * self.phase_errors.flatten())[:, None]
A_gain = self.gain_errors.flatten()[:, None]
DA *= A_phase
DA *= A_gain
return A, DA
def rotate(
self,
target_normal: np.ndarray,
current_normal: Optional[np.ndarray] = None,
center: Optional[np.ndarray] = None
) -> Tuple[BaseArray, np.ndarray]:
"""
Rotate a 1D array to align its normal vector with a target direction.
The array is assumed to be initially aligned along the y-axis (in 2D) or y-direction
in the xOy plane (in 3D). The method performs a rigid-body rotation to align the
array's normal with the specified `target_normal`.
The rotation is performed in three steps:
1. Translate the centroid of the array to the origin.
2. Apply the rotation based on the angle between current and target normal vectors.
3. Translate the array back to the specified or original centroid.
Args:
target_normal (np.ndarray): Target direction to align the array's normal to.
Must be a 2D or 3D vector of shape (2,) or (3,).
current_normal (Optional[np.ndarray]): Original normal direction of the array.
Defaults to [1, 0] in 2D and [0, 0, 1] in 3D.
center (Optional[np.ndarray]): Rotation center, defaults to the array's centroid.
Must be the same dimension as `target_normal`.
Returns:
Tuple[BaseArray, np.ndarray]:
- A new rotated array object (Array2D or Array3D) with updated element positions.
- The rotation matrix used in the transformation (2x2 or 3x3).
Raises:
ValueError: If `target_normal` is not 2D or 3D.
ValueError: If `center` is provided but its dimension does not match `target_normal`.
"""
target_normal = np.reshape(np.array(target_normal, dtype=float), (-1, 1))
dim = target_normal.shape[0]
if dim not in (2, 3):
raise ValueError(f"Target normal must be a 2D or 3D vector, got shape {target_normal.shape}.")
positions = self.positions
elem_num = positions.shape[1]
# Determine rotation center
if center is None:
center = np.mean(positions, axis=1)
center = np.reshape(np.array(center, dtype=float), (-1, 1))
if center.shape[0] != dim:
raise ValueError(
f"Center and target_normal must have the same dimension. "
f"Got center dim = {center.shape[0]}, target_normal dim = {dim}."
)
# Centered positions aligned along y-axis
if dim == 2:
# Arrange elements along Y-axis and center at origin
centered_pos = np.vstack((np.zeros((1, elem_num)), positions - center))
# Get normal vector
if current_normal is None:
current_normal = np.array([1.0, 0.0]) # assume
current_normal = current_normal / np.linalg.norm(current_normal)
target_normal = target_normal / np.linalg.norm(target_normal)
# Get the angle between current and target normal vector
angle = np.arctan2(target_normal[1], target_normal[0]) - np.arctan2(current_normal[1], current_normal[0])
rot_mat = np.array([
[np.cos(angle), -np.sin(angle)],
[np.sin(angle), np.cos(angle)]
])
positions_rot = rot_mat @ centered_pos + center[:, np.newaxis]
return Array2D(
positions = positions_rot,
type = self.type,
spacing = self.spacing,
indices = self.indices,
position_errors = self.position_errors,
gain_errors = self.gain_errors,
phase_errors = self.phase_errors
), rot_mat
elif dim == 3:
# Arrange elements along xOy plane and center at origin
centered_pos = np.vstack((
np.zeros((1, elem_num)),
positions - center,
np.zeros((1, elem_num))
))
# Get normal vector
if current_normal is None:
current_normal = np.array([0, 0, 1])
current_normal = current_normal / np.linalg.norm(current_normal)
target_normal = target_normal / np.linalg.norm(target_normal)
# Rotate to a 3D direction using axis-angle method
axis = np.cross(current_normal, target_normal)
norm_axis = np.linalg.norm(axis)
if norm_axis < 1e-8:
# If vectors are parallel or anti-parallel
if np.dot(current_normal, target_normal) > 0:
rot_mat = np.eye(3) # No rotation
else:
rot_mat = np.diag([-1, -1, 1]) # 180-degree rotation
else:
axis = axis / norm_axis
angle = np.arccos(np.dot(current_normal, target_normal))
K = np.array([
[0, -axis[2], axis[1]],
[axis[2], 0, -axis[0]],
[-axis[1], axis[0], 0]
])
rot_mat = np.eye(3) + np.sin(angle) * K + (1 - np.cos(angle)) * K @ K
positions_rot = rot_mat @ centered_pos + center[:, np.newaxis]
return Array3D(
positions = positions_rot,
type = self.type,
spacing = self.spacing,
indices = self.indices,
position_errors = self.position_errors,
gain_errors = self.gain_errors,
phase_errors = self.phase_errors
), rot_mat
else:
raise ValueError("Invalid dimension")
class Array2D(BaseArray):
"""
Two-dimensional planar array.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
if self.dim != 2:
raise ValueError(f"Array2D requires 2D array. Got dimension = {self.dim}.")
def steering_matrix(
self,
wavelength: float,
directions: np.ndarray,
) -> np.ndarray:
"""Compute the steering matrix of 2D array.
Args:
wavelength (float): Signal wavelength.
directions (np.ndarray): Input directions, must be in one of the following formats:
- shape (D,) or (1, D): D broadside angles (in radians), valid for 1D array
- shape (2, D): D pairs of azimuth and elevation angles (in radians), used for projection to broadside
Returns:
A (np.ndarray): Steering matrix of shape (M, D)
Raises:
ValueError: If dimension of directions is not 1 or 2.
"""
directions = np.array(directions)
if directions.ndim == 2:
if directions.shape[0] != 2:
raise ValueError(
f"If directions is 2D, its shape must be (2, D), "
f"but got shape {directions.shape}."
)
az, el = directions[0, :], directions[1, :]
cos_el = np.cos(el)
x_proj = cos_el * np.sin(az)
y_proj = cos_el * np.cos(az)
else:
raise ValueError(
f"Invalid shape for directions ({directions.ndim}). "
"Expected shape (D,) or (1, D) for 1D broadside angles or shape (2, D) for azimuth/elevation."
)
positions = self.positions
# Compute phase shifts
k = 2 * np.pi / wavelength
phase = k * (positions[0, :, None] * x_proj + positions[1, :, None] * y_proj)
A = np.exp(1j * phase)
# apply gain/phase errors
A *= np.exp(1j * self.phase_errors.flatten())
A *= self.gain_errors.flatten()
return A
def rotate_to(self, target_dim: int, rotation_matrix: np.ndarray) -> BaseArray:
assert target_dim == 3 and rotation_matrix.shape == (3, 2)
new_positions = self.element_positions @ rotation_matrix.T
return Array3D(...)
class Array3D(BaseArray):
"""
Three-dimensional volumetric array.
"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
assert self.dim == 3, "Array3D requires 3D element positions."
def steering_matrix(self, wavelength: float, direction: np.ndarray) -> np.ndarray:
"""
Compute 3D steering vector.
Args:
wavelength (float): Wavelength.
direction (np.ndarray): 3D unit vector.
Returns:
np.ndarray: Complex steering vector.
"""
k = 2 * np.pi / wavelength
return np.exp(1j * k * self.element_positions @ direction)
examples/ex1_ula.py
+10
View File
@@ -0,0 +1,10 @@
import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import numpy as np
from array_design.base_array import *
position = np.linspace(-2,2,5)
array_ula = Array1D(position, ArrayType.ULA)
requirements.txt
+5
View File
@@ -0,0 +1,5 @@
# requirements.txt
numpy>=1.23
scipy>=1.10
matplotlib>=3.7
pytest>=7.0