kaira.constraints.PerAntennaPowerConstraint

Inheritance diagram of PerAntennaPowerConstraint — Inheritance diagram for PerAntennaPowerConstraint

class kaira.constraints.PerAntennaPowerConstraint(power_budget: Tensor | None = None, uniform_power: float | None = None, *args, **kwargs)[source]

Bases: BaseConstraint

Distributes power budget across multiple antennas to ensure per-antenna power limits.

Ensures each antenna in a multi-antenna system (such as MIMO) adheres to its specific power budget. This constraint is crucial for systems where each antenna has its own power amplifier with individual power limitations [Yu and Lan, 2007] [Wunder et al., 2013].

Per-antenna power constraints are often more practical than sum-power constraints in real MIMO systems, as discussed in [Christopoulos et al., 2014] and [Yu and Lan, 2007].

The constraint can be configured either with individual power budgets for each antenna or with a uniform power value across all antennas.

power_budget

Power budget tensor for each antenna

Type:: torch.Tensor, optional

uniform_power

Uniform power level for all antennas

Type:: float, optional

Methods

`__init__`	Initialize the per-antenna power constraint.
`forward`	Apply per-antenna power constraint.
`get_dimensions`	Helper method to get all dimensions except batch for calculating norms/means.

__init__(power_budget: Tensor | None = None, uniform_power: float | None = None, *args, **kwargs) → None[source]

Initialize the per-antenna power constraint.

Parameters:

power_budget (torch.Tensor, optional) – Power budget for each antenna. Shape should be [num_antennas]. Mutually exclusive with uniform_power.
uniform_power (float, optional) – Uniform power value to apply across all antennas. Mutually exclusive with power_budget.
*args – Variable length argument list.
**kwargs – Arbitrary keyword arguments.

Raises:

AssertionError – If neither power_budget nor uniform_power is provided

Note

Either power_budget or uniform_power must be provided, but not both. If power_budget is provided, its length must match the number of antennas in the input signal.

forward(x: Tensor, *args, **kwargs) → Tensor[source]

Apply per-antenna power constraint.

Scales the signal from each antenna independently to meet its power budget. The second dimension of the input tensor is assumed to be the antenna dimension.

Parameters:

x (torch.Tensor) – Input tensor with shape [batch_size, num_antennas, …]. The second dimension must correspond to different antennas.
*args – Variable length argument list.
**kwargs – Arbitrary keyword arguments.

Returns:

Power-constrained signal with the same shape as input, where: each antenna’s signal has been scaled to meet its power budget

Return type:

torch.Tensor

Note

Power is calculated by averaging the squared magnitude across all dimensions except batch and antenna dimensions.

static get_dimensions(x: Tensor, exclude_batch: bool = True) → Tuple[int, ...]

Helper method to get all dimensions except batch for calculating norms/means.

Utility function to generate dimension indices for reduction operations like mean or norm. Typically used to calculate signal properties across all dimensions except the batch dimension.

Parameters:

x (torch.Tensor) – Input tensor
exclude_batch (bool, optional) – Whether to exclude the batch dimension (first dimension). Defaults to True.

Returns:

Dimensions to use for reduction operations (e.g., mean, norm)

Return type:

Tuple[int, …]

Example

>>> x = torch.randn(32, 4, 128)  # [batch, antennas, time]
>>> dims = BaseConstraint.get_dimensions(x)
>>> # dims will be (1, 2) for summing across antennas and time