From 8986a18095e70033fff4fa76a6e0ba6c8b94e35b Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Fri, 30 Jan 2026 15:29:48 -0300
Subject: [PATCH 01/10] wip: optimized M.2101 array antenna implementation

---
 sharc/antenna/antenna_array.py | 318 +++++++++++++++++++++++++++++
 tests/test_antenna_array.py    | 361 +++++++++++++++++++++++++++++++++
 2 files changed, 679 insertions(+)
 create mode 100644 sharc/antenna/antenna_array.py
 create mode 100644 tests/test_antenna_array.py

diff --git a/sharc/antenna/antenna_array.py b/sharc/antenna/antenna_array.py
new file mode 100644
index 00000000..582633be
--- /dev/null
+++ b/sharc/antenna/antenna_array.py
@@ -0,0 +1,318 @@
+"""
+This antenna was created after the already existing antenna_beamforming_imt
+since that implementation is too slow for use with a lot of stations.
+
+This is supposed to be a faster implementation, and substitute the previous one
+in a future release
+"""
+
+from sharc.antenna.antenna import Antenna
+from sharc.parameters.imt.parameters_antenna_imt import ParametersAntennaImt
+from sharc.support.geometry import RigidTransform
+
+import numpy as np
+
+
+class AntennaArray(Antenna):
+    # par: ParametersAntennaImt
+
+    def __init__(
+        self,
+        par: ParametersAntennaImt,
+        # global2local_transform: RigidTransform
+    ):
+        super().__init__()
+        self.par = par
+
+    def calculate_gain(self, *args, **kwargs) -> np.array:
+        """
+        Calculates the antenan gain.
+        """
+        phi_vec = np.atleast_1d(kwargs["phi_vec"])
+        theta_vec = np.atleast_1d(kwargs["theta_vec"])
+        co_channel = kwargs.get("co_channel", True)
+        adj_antenna_model = (
+            self.par.adjacent_antenna_model == "SINGLE_ELEMENT"
+            and not co_channel
+        )
+        if "beams_l" in kwargs.keys():
+            beam_idxs = np.asarray(kwargs["beams_l"], dtype=int)
+        else:
+            beam_idxs = np.arange(len(phi_vec))
+
+        assert phi_vec.shape == (len(phi_vec),)
+        assert theta_vec.shape == (len(theta_vec),)
+        phi_vec, theta_vec = self._to_local_coord(
+            phi_vec,
+            theta_vec,
+        )
+
+        el_g = self._element_gain(
+            phi_vec, theta_vec,
+        )
+        assert el_g.shape == theta_vec.shape
+
+        if adj_antenna_model:
+            return el_g
+
+        ar_g = self._array_gain(
+            phi_vec, theta_vec, beam_idxs
+        )
+        assert ar_g.shape == theta_vec.shape
+
+        return ar_g + el_g
+
+    def _element_gain(
+        self, phi: np.ndarray, theta: np.ndarray
+    ):
+        return self._element_gain_dispatch(
+            self.par, phi, theta,
+        )
+
+    def _array_gain(
+        self,
+        phi: np.ndarray, theta: np.ndarray,
+        beam_idxs: np.ndarray
+    ):
+        v_vec = self._super_position_vector(
+            phi, theta,
+            self.par.n_rows, self.par.n_columns,
+            self.par.element_vert_spacing,
+            self.par.element_horiz_spacing,
+        )
+        if len(self.beams_list) == 0:
+            beam_phi, beam_theta = phi, theta
+        else:
+            beam_phi, beam_theta = np.array(self.beams_list).T
+
+        beam_etilt = beam_theta - 90.
+        beams_w_vec = self._weight_vector(
+            beam_phi, beam_etilt,
+            self.par.n_rows, self.par.n_columns,
+            self.par.element_vert_spacing,
+            self.par.element_horiz_spacing,
+        )
+        w_vec = beams_w_vec[beam_idxs]
+
+        g = 10 * np.log10(
+            abs(
+                np.sum(v_vec * w_vec, axis=(1, 2))
+            )**2
+        )
+
+        return g
+
+    @staticmethod
+    def _weight_vector(
+        phi_tilt: np.ndarray,
+        theta_tilt: np.ndarray,
+        n_rows: int, n_cols: int,
+        dv: float, dh: float,
+    ) -> np.array:
+        """
+        Calculates super position vector.
+        Angles are in the local coordinate system.
+
+        Parameters
+        ----------
+            phi_tilt (float): electrical horizontal steering [degrees]
+            theta_tilt (float): electrical down-tilt steering [degrees]
+
+        Returns
+        -------
+            w_vec (np.array): weighting vector
+        """
+        # shape (Na, 1, 1)
+        r_phi = np.atleast_1d(
+            np.deg2rad(phi_tilt)
+        )[:, np.newaxis, np.newaxis]
+        r_theta = np.atleast_1d(
+            np.deg2rad(theta_tilt)
+        )[:, np.newaxis, np.newaxis]
+
+        # shape (1, Nr, 1)
+        n = np.arange(n_rows)[np.newaxis, :, np.newaxis] + 1
+        # shape (1, 1, Nc)
+        m = np.arange(n_cols)[np.newaxis, np.newaxis, :] + 1
+
+        exp_arg = (n - 1) * dv * np.sin(r_theta) - \
+                  (m - 1) * dh * np.cos(r_theta) * np.sin(r_phi)
+
+        w_vec = (1 / np.sqrt(n_rows * n_cols)) *\
+            np.exp(2 * np.pi * 1.0j * exp_arg)
+
+        # shape (Na, Nr, Nc)
+        return w_vec
+
+    @staticmethod
+    def _super_position_vector(
+        phi: float, theta: float,
+        n_rows: int, n_cols: int,
+        dv: float, dh: float,
+    ) -> np.array:
+        """
+        Calculates super position vector.
+        Angles are in the local coordinate system.
+
+        Parameters
+        ----------
+            theta (float): elevation angle [degrees]
+            phi (float): azimuth angle [degrees]
+
+        Returns
+        -------
+            v_vec (np.array): superposition vector
+        """
+        # shape (Na, 1, 1)
+        r_phi = np.atleast_1d(
+            np.deg2rad(phi)
+        )[:, np.newaxis, np.newaxis]
+        r_theta = np.atleast_1d(
+            np.deg2rad(theta)
+        )[:, np.newaxis, np.newaxis]
+
+        # shape (1, Nr, 1)
+        n = np.arange(n_rows)[np.newaxis, :, np.newaxis] + 1
+        # shape (1, 1, Nc)
+        m = np.arange(n_cols)[np.newaxis, np.newaxis, :] + 1
+
+        exp_arg = (n - 1) * dv * np.cos(r_theta) + \
+                  (m - 1) * dh * np.sin(r_theta) * np.sin(r_phi)
+
+        v_vec = np.exp(2 * np.pi * 1.0j * exp_arg)
+
+        # shape (Na, Nr, Nc)
+        return v_vec
+
+    @staticmethod
+    def _element_gain_dispatch(par: ParametersAntennaImt, phi, theta):
+        if par.element_pattern == "M2101":
+            return AntennaArray._calculate_m2101_element_gain(
+                phi, theta,
+                par.element_phi_3db, par.element_theta_3db,
+                par.element_max_g, par.element_sla_v, par.element_am,
+                par.multiplication_factor,
+            )
+        else:
+            raise NotImplementedError(
+                "No implementation done for element_pattern"
+                f"='{par.element_pattern}'"
+            )
+
+    @staticmethod
+    def _calculate_m2101_element_gain(
+        phi: np.ndarray, theta: np.ndarray,
+        phi_3db: np.ndarray, theta_3db: np.ndarray,
+        g_max: np.ndarray, sla_v: np.ndarray, am: np.ndarray,
+        multiplication_factor: np.ndarray = 12
+    ):
+        """Calculates and returns element gain as described in M.2101
+        """
+        g_horizontal = -1.0 * np.minimum(
+            multiplication_factor * (phi / phi_3db)**2, am
+        )
+        g_vertical = -1.0 * np.minimum(
+            multiplication_factor * ((theta-90.) / theta_3db)**2, sla_v
+        )
+
+        att = -1.0 * (
+            g_horizontal +
+            g_vertical
+        )
+
+        return g_max - np.minimum(att, am)
+
+    def _to_local_coord(self, phi, theta):
+        return np.array(phi), np.array(theta)
+
+    def add_beam(self, phi_etilt: float, theta_etilt: float):
+        """
+        Add new beam to antenna.
+        Does not receive angles in local coordinate system.
+        Theta taken with z axis as reference.
+
+        Parameters
+        ----------
+            phi_etilt (float): azimuth electrical tilt angle [degrees]
+            theta_etilt (float): elevation electrical tilt angle [degrees]
+        """
+        # phi_etilt, theta_etilt = np.atleast_1d(phi_etilt), np.atleast_1d(theta_etilt)
+        phi, theta = self._to_local_coord(phi_etilt, theta_etilt)
+        self.beams_list.append(
+            (np.ndarray.item(phi), np.ndarray.item(theta)),
+        )
+
+
+if __name__ == "__main__":
+    antenna_params = ParametersAntennaImt()
+    antenna_params.adjacent_antenna_model = "SINGLE_ELEMENT"
+    antenna_params.normalization = False
+    antenna_params.minimum_array_gain = -200
+
+    antenna_params.element_pattern = "M2101"
+    antenna_params.element_max_g = 6.5
+    antenna_params.element_phi_3db = 65
+    antenna_params.element_theta_3db = 90
+    antenna_params.element_am = 30
+    antenna_params.element_sla_v = 30
+    antenna_params.n_rows = 8
+    antenna_params.n_columns = 8
+    antenna_params.element_horiz_spacing = 0.5
+    antenna_params.element_vert_spacing = 0.5
+    antenna_params.multiplication_factor = 12
+
+    par = antenna_params.get_antenna_parameters()
+    antenna = AntennaArray(par)
+
+    antenna.add_beam(np.array(0.), np.array(90.))
+
+    phi_scan = np.linspace(-180., 180., num=360)
+    theta = np.zeros_like(phi_scan) + 90.
+
+    # gain = antenna._element_gain(
+    #     phi_scan,
+    #     theta,
+    # )
+    gain = antenna.calculate_gain(
+        phi_vec=phi_scan,
+        theta_vec=theta,
+        beams_l=np.zeros_like(phi_scan),
+    )
+
+    import matplotlib.pyplot as plt
+
+    fig = plt.figure(figsize=(15, 5), facecolor='w', edgecolor='k')
+    ax1 = fig.add_subplot(121)
+
+    ax1.plot(phi_scan, gain)
+    top_y_lim = np.ceil(np.max(gain) / 10) * 10
+    ax1.set_xlim(-180, 180)
+    ax1.set_ylim(top_y_lim - 60, top_y_lim)
+    ax1.grid(True)
+    ax1.set_xlabel(r"$\varphi$ [deg]")
+    ax1.set_ylabel("Gain [dBi]")
+
+    theta_scan = np.linspace(0., 180., num=360)
+    phi = np.zeros_like(theta_scan)
+
+    # gain = antenna._element_gain(
+    #     phi,
+    #     theta_scan,
+    # )
+    gain = antenna.calculate_gain(
+        phi_vec=phi,
+        theta_vec=theta_scan,
+        beams_l=np.zeros_like(theta_scan),
+    )
+    # fig = plt.figure(figsize=(15, 5), facecolor='w', edgecolor='k')
+    ax2 = fig.add_subplot(122, sharey=ax1)
+
+    ax2.plot(theta_scan, gain)
+    top_y_lim = np.ceil(np.max(gain) / 10) * 10
+    ax2.set_xlim(0, 180.)
+    ax2.set_ylim(top_y_lim - 60, top_y_lim)
+    ax2.grid(True)
+    ax2.set_xlabel(r"$\vartheta$ [deg]")
+    ax2.set_ylabel("Gain [dBi]")
+
+    plt.show()
diff --git a/tests/test_antenna_array.py b/tests/test_antenna_array.py
new file mode 100644
index 00000000..adaf1291
--- /dev/null
+++ b/tests/test_antenna_array.py
@@ -0,0 +1,361 @@
+# -*- coding: utf-8 -*-
+"""
+Created on Sat Apr 15 15:36:22 2017
+
+@author: Calil
+"""
+
+import unittest
+import numpy as np
+import numpy.testing as npt
+
+from sharc.antenna.antenna_array import AntennaArray
+from sharc.parameters.imt.parameters_antenna_imt import ParametersAntennaImt
+
+
+class AntennaArrayTest(unittest.TestCase):
+    """Unit tests for the AntennaArray class."""
+
+    def setUp(self):
+        """Set up test fixtures for AntennaArray tests."""
+        # Array parameters
+        self.bs_param = ParametersAntennaImt()
+        self.ue_param = ParametersAntennaImt()
+
+        # NOTE: not implemented:
+        self.bs_param.adjacent_antenna_model = "SINGLE_ELEMENT"
+        self.bs_param.normalization = False
+        self.bs_param.normalization_file = None
+        self.ue_param.adjacent_antenna_model = "SINGLE_ELEMENT"
+        self.ue_param.normalization = False
+        self.ue_param.normalization_file = None
+
+        self.bs_param.element_pattern = "M2101"
+        self.bs_param.minimum_array_gain = -200
+        self.bs_param.downtilt = 0
+        self.bs_param.element_max_g = 5
+        self.bs_param.element_phi_3db = 80
+        self.bs_param.element_theta_3db = 60
+        self.bs_param.element_am = 30
+        self.bs_param.element_sla_v = 30
+        self.bs_param.n_rows = 16
+        self.bs_param.n_columns = 16
+        self.bs_param.element_horiz_spacing = 1
+        self.bs_param.element_vert_spacing = 1
+        self.bs_param.multiplication_factor = 12
+
+        self.ue_param.element_pattern = "M2101"
+        self.ue_param.minimum_array_gain = -200
+        self.ue_param.element_max_g = 10
+        self.ue_param.element_phi_3db = 75
+        self.ue_param.element_theta_3db = 65
+        self.ue_param.element_am = 25
+        self.ue_param.element_sla_v = 35
+        self.ue_param.n_rows = 2
+        self.ue_param.n_columns = 2
+        self.ue_param.element_horiz_spacing = 0.5
+        self.ue_param.element_vert_spacing = 0.5
+        self.ue_param.multiplication_factor = 12
+        # Create antenna objects
+        par = self.bs_param.get_antenna_parameters()
+        self.antenna1 = AntennaArray(par)
+        par = self.ue_param.get_antenna_parameters()
+        self.antenna2 = AntennaArray(par)
+
+    def test_element_gain(self):
+        """Testing element gain calculations"""
+        # self.antenna1._calculate_m2101_element_gain()
+        """Test M.2101 horizontal pattern calculation for various phi values."""
+        # phi = 0 results in zero gain
+        phi = np.array([0., 120., 150.])
+        theta = np.zeros_like(phi) + 90.
+        h_att = self.antenna1._element_gain(phi, theta)
+        npt.assert_equal(
+            h_att,
+            self.antenna1.par.element_max_g - np.array([0.0, 27.0, 30.0])
+        )
+
+        """Test M.2101 vertical pattern calculation for various theta values."""
+        theta = np.array([90, 180, 210])
+        phi = np.zeros_like(theta)
+        v_att = self.antenna1._element_gain(phi, theta)
+        npt.assert_equal(
+            v_att,
+            self.antenna1.par.element_max_g - np.array([0.0, 27.0, 30.0])
+        )
+
+        """Test element pattern calculation for various phi/theta values."""
+        phi = np.array([0, 80, 150])
+        theta = np.array([90, 150, 210])
+        e_gain = self.antenna1._element_gain(phi, theta)
+        npt.assert_equal(
+            e_gain,
+            np.array([5.0, -19.0, -25.0])
+        )
+
+    def test_weight_vector(self):
+        """Test calculation of the weight vector for beamforming."""
+        # Error margin
+        eps = 1e-5
+
+        acc_phi_scan = []
+        acc_theta_tilt = []
+        acc_w_vec = []
+        # Test 1
+        phi_scan = 0
+        theta_tilt = 0
+        w_vec = self.antenna2._weight_vector(
+            phi_scan, theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi_scan.append(phi_scan)
+        acc_theta_tilt.append(theta_tilt)
+        acc_w_vec.append(w_vec)
+        expected_w_vec = np.array([[[0.5, 0.5], [0.5, 0.5]]])
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+        # Test 2
+        phi_scan = 90
+        theta_tilt = 90
+        w_vec = self.antenna2._weight_vector(
+            phi_scan, theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi_scan.append(phi_scan)
+        acc_theta_tilt.append(theta_tilt)
+        acc_w_vec.append(w_vec)
+        expected_w_vec = np.array([[[0.5, 0.5], [-0.5, -0.5]]])
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+        # Test 3
+        phi_scan = 45
+        theta_tilt = 45
+        w_vec = self.antenna2._weight_vector(
+            phi_scan, theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi_scan.append(phi_scan)
+        acc_theta_tilt.append(theta_tilt)
+        acc_w_vec.append(w_vec)
+        expected_w_vec = np.array([[
+            [0.5 + 0.0j, 0.0 - 0.5j],
+            [-0.3028499 + 0.3978466j, 0.3978466 + 0.3028499j],
+        ]])
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+        # Test 4
+        phi_scan = 0
+        theta_tilt = 90
+        w_vec = self.antenna2._weight_vector(
+            phi_scan, theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi_scan.append(phi_scan)
+        acc_theta_tilt.append(theta_tilt)
+        acc_w_vec.append(w_vec)
+        expected_w_vec = np.array([[[0.5, 0.5], [-0.5, -0.5]]])
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+        # Test 5
+        phi_scan = 45
+        theta_tilt = 30
+        w_vec = self.antenna2._weight_vector(
+            phi_scan, theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi_scan.append(phi_scan)
+        acc_theta_tilt.append(theta_tilt)
+        acc_w_vec.append(w_vec)
+        expected_w_vec = np.array([[
+            [0.5 + 0.0j, -0.172870 - 0.469169j],
+            [0.0 + 0.5j, 0.469165 - 0.172870j],
+        ]])
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+        acc_phi_scan = np.array(acc_phi_scan)
+        acc_theta_tilt = np.array(acc_theta_tilt)
+        acc_w_vec = np.array(acc_w_vec)
+        w_vec = self.antenna2._weight_vector(
+            acc_phi_scan, acc_theta_tilt,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        expected_w_vec = np.squeeze(acc_w_vec, axis=1)
+        npt.assert_allclose(
+            w_vec,
+            expected_w_vec, rtol=eps,
+        )
+
+    def test_super_position_vector(self):
+        """Test calculation of the superposition vector."""
+        # Error margin
+        eps = 1e-5
+
+        acc_phi = []
+        acc_theta = []
+        acc_v_vec = []
+
+        # Test 1
+        phi = 0
+        theta = 0
+        v_vec = self.antenna2._super_position_vector(
+            phi, theta,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi.append(phi)
+        acc_theta.append(theta)
+        acc_v_vec.append(v_vec)
+        expected_v_vec = np.array([[[1.0, 1.0], [-1.0, -1.0]]])
+        npt.assert_allclose(
+            v_vec,
+            expected_v_vec, rtol=eps,
+        )
+
+        # Test 2
+        phi = 90
+        theta = 90
+        v_vec = self.antenna2._super_position_vector(
+            phi, theta,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi.append(phi)
+        acc_theta.append(theta)
+        acc_v_vec.append(v_vec)
+        expected_v_vec = np.array([[[1.0, -1.0], [1.0, -1.0]]])
+        npt.assert_allclose(
+            v_vec,
+            expected_v_vec, rtol=eps,
+        )
+
+        # Test 3
+        phi = 45
+        theta = 45
+        v_vec = self.antenna2._super_position_vector(
+            phi, theta,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi.append(phi)
+        acc_theta.append(theta)
+        acc_v_vec.append(v_vec)
+        expected_v_vec = np.array([[
+            [1.0 + 0.0j, 0.0 + 1.0j],
+            [-0.6056998 + 0.7956932j, -0.7956932 - 0.6056998j],
+        ]])
+        npt.assert_allclose(
+            v_vec,
+            expected_v_vec, rtol=eps,
+        )
+
+        # Test 4
+        phi = 60
+        theta = 90
+        v_vec = self.antenna2._super_position_vector(
+            phi, theta,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        acc_phi.append(phi)
+        acc_theta.append(theta)
+        acc_v_vec.append(v_vec)
+        expected_v_vec = np.array([[
+            [1.0 + 0.0j, -0.912724 + 0.408576j],
+            [1.0 + 0.0j, -0.912724 + 0.408576j],
+        ]])
+        npt.assert_allclose(
+            v_vec,
+            expected_v_vec, rtol=eps,
+        )
+
+        acc_phi = np.array(acc_phi)
+        acc_theta = np.array(acc_theta)
+        expected_v_vec = np.squeeze(acc_v_vec, axis=1)
+        v_vec = self.antenna2._super_position_vector(
+            acc_phi, acc_theta,
+            self.antenna2.par.n_rows, self.antenna2.par.n_columns,
+            self.antenna2.par.element_vert_spacing,
+            self.antenna2.par.element_horiz_spacing,
+        )
+        npt.assert_allclose(
+            v_vec,
+            expected_v_vec, rtol=eps,
+        )
+
+    def test_calculate_gain(self):
+        """Test calculation of antenna gain for given phi/theta vectors."""
+        # Error margin and antenna
+        eps = 1e-4
+        par = self.bs_param.get_antenna_parameters()
+        self.antenna1 = AntennaArray(par)
+        par = self.ue_param.get_antenna_parameters()
+        self.antenna2 = AntennaArray(par)
+
+        # Test 1
+        phi_vec = np.array([45.0, 32.5])
+        theta_vec = np.array([45.0, 115.2])
+        gains = self.antenna2.calculate_gain(
+            phi_vec=phi_vec, theta_vec=theta_vec,
+        )
+        npt.assert_allclose(gains, np.array([5.9491, 11.9636]), atol=eps)
+
+        # Test 2
+        phi = 0.0
+        theta = 60.0
+        phi_scan = 45
+        theta_tilt = 180
+        self.antenna2.add_beam(phi_scan, theta_tilt)
+        beams_l = np.zeros_like(phi, dtype=int)
+        gains = self.antenna2.calculate_gain(
+            phi_vec=phi, theta_vec=theta,
+            beams_l=beams_l,
+        )
+        npt.assert_allclose(gains, np.array([10.454087]), atol=eps)
+
+        # Test 3
+        phi = 40
+        theta = 100
+        gains = self.antenna1.calculate_gain(
+            phi_vec=phi, theta_vec=theta,
+            co_channel=False,
+        )
+        npt.assert_allclose(gains, np.array([1.6667]), atol=eps)
+
+
+if __name__ == '__main__':
+    unittest.main()
+#
+#    suite = unittest.TestSuite()
+#    suite.addTest(AntennaArrayTest('test_calculate_gain'))
+#    unittest.TextTestRunner().run(suite)

From 86f1f72d5d6d25faa6be95fb87f59a3f60897376 Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Fri, 30 Jan 2026 15:42:25 -0300
Subject: [PATCH 02/10] update!: simulator geometry separate concerns refactor
 and better API

This commit adds RigidTransform, ReferenceFrame and ENUReferenceFrame
for preparation for usage with other reference frames. SimulatorGeometry
was refactored and its API was updated for better maintainability.
---
 sharc/support/geometry.py                     | 439 +++++++++++------
 .../topology_spherical_sampling_from_grid.py  |  45 +-
 tests/test_geometry.py                        | 463 ++++++++++++++++--
 3 files changed, 730 insertions(+), 217 deletions(-)

diff --git a/sharc/support/geometry.py b/sharc/support/geometry.py
index 6cf9e51f..94a93712 100644
--- a/sharc/support/geometry.py
+++ b/sharc/support/geometry.py
@@ -5,6 +5,7 @@
 from sharc.support.sharc_geom import cartesian_to_polar, polar_to_cartesian
 import scipy
 import numpy as np
+from dataclasses import dataclass
 from abc import ABC
 
 
@@ -303,10 +304,181 @@ def get_off_axis_angle(self, other: "GlobalGeometry") -> np.array:
         return phi_deg
 
 
+@dataclass(frozen=True)
+class RigidTransform:
+    """
+    Defines a transformation of the type
+        A(X) = rot @ X + t
+    where it applies rotation to X and then translates the result.
+    """
+    rot: scipy.spatial.transform.Rotation  # N rotations
+    t: np.ndarray                          # (N, 3)
+
+    def __post_init__(self):
+        if not isinstance(self.rot, scipy.spatial.transform.Rotation):
+            raise ValueError("rot must be a scipy Rotation")
+
+        n_rot = len(self.rot)
+        n_t = self.t.shape[0]
+
+        if self.t.shape != (n_t, 3):
+            raise ValueError(f"Invalid transform shape t={n_t}")
+
+        # if some is equal to one, broadcasting still works
+        if not (n_rot == n_t or n_rot == 1 or n_t == 1):
+            raise ValueError(
+                f"Incompatible RigidTransform shapes: rot={n_rot}, t={n_t}"
+            )
+        self.t.flags.writeable = False
+
+    @property
+    def N(self):
+        n_rot = len(self.rot)
+        n_t = self.t.shape[0]
+        return max(n_rot, n_t)
+
+    def inv(self) -> "RigidTransform":
+        rot_inv = self.rot.inv()
+        t_inv = -rot_inv.apply(self.t)
+        return RigidTransform(rot_inv, t_inv)
+
+    def and_then(self, other: "RigidTransform") -> "RigidTransform":
+        """Apply self, then other (composition: other @ self)."""
+        return RigidTransform(
+            rot=other.rot * self.rot,
+            t=other.rot.apply(self.t) + other.t
+        )
+
+    def take(self, idx: int) -> "RigidTransform":
+        # NOTE: we slice instead of taking indice to maintain
+        # array shape structure for functional broadcasting
+        # and batch computing
+        return RigidTransform(
+            rot=self.rot[idx:idx+1],
+            t=self.t[idx:idx+1],
+        )
+
+    def apply_points(self, x: np.ndarray) -> np.ndarray:
+        """Applies rotations and translations to the points"""
+        return self.apply_vectors(x) + self.t
+
+    def apply_points_permutation(self, x: np.ndarray) -> np.ndarray:
+        """Applies rotation and translation to the point"""
+        t = self.t                     # (N,3)
+
+        return self.apply_vectors_permutation(x) + t[:, None, :]
+
+    def apply_vectors(self, v: np.ndarray) -> np.ndarray:
+        """Applies only rotations, considering the vectors as a pointing vec"""
+        v = np.atleast_2d(v)
+        assert v.ndim == 2 and v.shape[1] == 3
+        M = v.shape[0]
+        assert self.N == 1 or M == 1 or M == self.N
+
+        return self.rot.apply(v) + np.zeros_like(self.t)
+
+    def apply_vectors_permutation(self, v: np.ndarray) -> np.ndarray:
+        """Applies rotation and translation to the point"""
+        # TODO: change impl
+        v = np.atleast_2d(v)
+        assert v.ndim == 2 and v.shape[1] == 3
+
+        # (N,3,3)
+        R = self.rot.as_matrix()
+
+        # einsum: (N,3,3) . (M,3) -> (N,M,3)
+        return (
+            np.einsum("nij,mj->nmi", R, v)
+            + np.zeros_like(self.t)[:, None, :]
+        )
+
+
+class ReferenceFrame(ABC):
+    """
+    Defines a reference frame, which generates a rigid transform
+    between ecef and its own local coordinate system.
+    It should be noted that, while this may define axis position,
+    the bearing considered (azimuth, elevation) for calculations
+    CANNOT USE NAVIGATION CONVENTION.
+    """
+    __slots__ = ('_from_ecef', '_to_ecef')
+
+    @property
+    def from_ecef(self) -> RigidTransform:
+        return self._from_ecef
+
+    @property
+    def to_ecef(self) -> RigidTransform:
+        return self._to_ecef
+
+
+class ENUReferenceFrame(ReferenceFrame):
+    """
+    Defines ENU reference frame.
+
+    NOTE: this does not change bearing convention to azimuth being
+    angular distance from North. Azimuth is still from x axis towards y axis.
+    """
+    __slots__ = ("_lat", "_lon", "_alt")
+
+    def __init__(
+        self,
+        *,
+        lat: np.ndarray,
+        lon: np.ndarray,
+        alt: np.ndarray,
+    ):
+        lat = np.atleast_1d(lat)
+        lon = np.atleast_1d(lon)
+        alt = np.atleast_1d(alt)
+        if lat.shape != lon.shape or lat.shape != alt.shape:
+            raise ValueError("lat, lon, alt must have identical shapes")
+
+        if lat.shape != (len(lat),):
+            raise ValueError(
+                "The lla shapes used must be one dimensional: (N,), but is"
+                f" {lat.shape}"
+            )
+
+        self._lat = np.copy(lat)
+        self._lon = np.copy(lon)
+        self._alt = np.copy(alt)
+
+        self._from_ecef = self._compute_to_local()
+        self._to_ecef = self._from_ecef.inv()
+
+        # freeze frame
+        self._lat.flags.writeable = False
+        self._lon.flags.writeable = False
+        self._alt.flags.writeable = False
+
+    def _compute_to_local(self):
+        lat = self._lat
+        lon = self._lon
+        alt = self._alt
+        N = lat.shape[0]
+
+        rotation_around_z = -lon - 90
+        rotation_around_x = lat - 90
+
+        ecef2local_rot = scipy.spatial.transform.Rotation.from_euler(
+            'zx',
+            np.stack([rotation_around_z, rotation_around_x], axis=-1),
+            degrees=True
+        )
+        ecef2local_translation = np.zeros((N, 3))
+        # NOTE: using constant earth radius works because of spherical Earth
+        ecef2local_translation[:, 2] = -(alt + EARTH_RADIUS_M)
+        ecef2local = RigidTransform(
+            ecef2local_rot, ecef2local_translation
+        )
+        return ecef2local
+
+
 @readonly_properties(
     "x_local", "y_local", "z_local",
     "pointn_azim_local", "pointn_elev_local",
-    "global_lla_reference", "local_lla_references"
+    "global_reference_frame", "local_reference_frame"
 )
 class SimulatorGeometry(GlobalGeometry):
     """
@@ -320,8 +492,8 @@ class SimulatorGeometry(GlobalGeometry):
     pointn_azim_local: np.ndarray  # (N,)
     pointn_elev_local: np.ndarray  # (N,)
 
-    local_lla_references: np.ndarray[np.ndarray[float]]  # (3, N)
-    global_lla_reference: tuple[float, float, float]
+    local_reference_frame: ReferenceFrame
+    global_reference_frame: ReferenceFrame
 
     uses_local_coords: bool
 
@@ -345,14 +517,14 @@ def setup(
         self,
         num_geometries,
         uses_local_coords=False,
-        global_cs: tuple[float, float, float] = None,
+        global_cs: ReferenceFrame = None,
     ):
         """
         Initializes variables based on number of geometries considered
         """
         super().setup(num_geometries)
 
-        self._global_lla_reference = global_cs
+        self._global_reference_frame = global_cs
         self.uses_local_coords = True
 
         if not uses_local_coords:
@@ -373,25 +545,16 @@ def setup(
         self._z_local = np.empty(num_geometries)
         self._pointn_azim_local = np.empty(num_geometries)
         self._pointn_elev_local = np.empty(num_geometries)
-        self._local_lla_references = np.empty((3, num_geometries))
+        self._local_reference_frame = None
 
-    def set_local_coord_sys(
-        self,
-        ref_lats,
-        ref_lons,
-        ref_alts,
+    def set_local_reference_frame(
+        self, frame: ReferenceFrame
     ):
         """
         Sets local coord system references and prepares
         global<->local coordinate transformation
         """
-        for r in [ref_lats, ref_lons, ref_alts]:
-            if len(r) != self.num_geometries:
-                raise ValueError(
-                    "Incongruent number of coordinate systems. "
-                    f"Passed {len(r)} but should have passed {self.num_geometries}"
-                )
-        self._local_lla_references = np.stack((ref_lats, ref_lons, ref_alts))
+        self._local_reference_frame = frame
 
         self._compute_global_local_transform()
 
@@ -453,7 +616,7 @@ def _compute_global_from_local(self):
 
         p_local = np.stack([self.x_local, self.y_local, self.z_local], axis=-1)
 
-        p_global = self._vec_local2global(p_local)
+        p_global = self._local2global_points(p_local)
 
         # Store results
         self._x_global = p_global[:, 0]
@@ -465,8 +628,8 @@ def _compute_global_from_local(self):
         point_local = np.stack(polar_to_cartesian(
             r, self.pointn_azim_local, self.pointn_elev_local), axis=-1)
 
-        point_global_x, point_global_y, point_global_z = self._vec_local2global(
-            point_local, translate=False
+        point_global_x, point_global_y, point_global_z = self._local2global_vectors(
+            point_local
         ).T
 
         _, global_azimuth, global_elevation = cartesian_to_polar(
@@ -475,42 +638,73 @@ def _compute_global_from_local(self):
         self._pointn_elev_global = global_elevation
         self._pointn_azim_global = global_azimuth
 
-    def _vec_local2global(
-        self,
-        p_local,
-        *,
-        translate=True,
-        permutate=False,
+    def _local2global_points(
+        self, p_local,
+    ):
+        """Receives points shaped as (N, 3) and applies local2global transform,
+        returning (N, 3), where N is the number of local coordinate systems.
+        """
+        return self.local2global.apply_points(p_local)
+
+    def _local2global_vectors(
+        self, p_local,
     ):
         """Receives a vector shaped as (N, 3) and applies local2global transform,
-            returning (N, 3).
-        If permutate == True, then may receive a vector (M, 3),
-            returning (N, M, 3)
+        returning (N, 3), where N is the number of local coordinate systems.
+        Does not apply translations, just rotations
+        """
+        return self.local2global.apply_vectors(p_local)
+
+    def _local2global_points_permutation(
+        self, p_local,
+    ):
+        """Receives points as (M, 3), returning (N, M, 3),
+        where N is the number of local coordinate systems
+        """
+        return self.local2global.apply_points_permutation(p_local)
+
+    def _local2global_vectors_permutation(
+        self, p_local,
+    ):
+        """Receives points as (M, 3), returning (N, M, 3),
+        where N is the number of local coordinate systems.
+        Does not apply translations, just rotations
+        """
+        return self.local2global.apply_vectors_permutation(p_local)
+
+    def _global2local_points(
+        self, p_local,
+    ):
+        """Receives points shaped as (N, 3) and applies global2local transform,
+        returning (N, 3), where N is the number of local coordinate systems.
+        """
+        return self.global2local.apply_points(p_local)
+
+    def _global2local_vectors(
+        self, p_local,
+    ):
+        """Receives a vector shaped as (N, 3) and applies global2local transform,
+        returning (N, 3), where N is the number of local coordinate systems.
+        Does not apply translations, just rotations
+        """
+        return self.global2local.apply_vectors(p_local)
+
+    def _global2local_points_permutation(
+        self, p_local,
+    ):
+        """Receives points as (M, 3), returning (N, M, 3),
+        where N is the number of local coordinate systems
+        """
+        return self.global2local.apply_points_permutation(p_local)
+
+    def _global2local_vectors_permutation(
+        self, p_local,
+    ):
+        """Receives points as (M, 3), returning (N, M, 3),
+        where N is the number of local coordinate systems.
+        Does not apply translations, just rotations
         """
-        N = self.num_geometries
-
-        local2ecef_t = self._local2ecef_transl_mtx if translate else 0
-        ecef2global_t = self._ecef2global_transl_mtx if translate else 0
-
-        if not permutate:
-            assert p_local.shape == (N, 3)
-            # remove ecef2local translation
-            p_local = p_local + local2ecef_t  # (N,3)
-            # rotate local2ecef2global
-            p_global = (self._local2global_rot_mtx @ p_local[..., None]).squeeze(-1)  # (N,3)
-            # add ecef2global translation
-            p_global = p_global + ecef2global_t  # (N,3)
-        else:
-            local2ecef_t_cast = local2ecef_t
-            ecef2global_t_cast = ecef2global_t
-            if translate:
-                local2ecef_t_cast = local2ecef_t[:, None, :]
-                ecef2global_t_cast = ecef2global_t[:, None, :]
-            p_local_exp = p_local[None, :, :] + local2ecef_t_cast  # (N,M,3)
-            p_global = (self._global2local_rot_mtx[:, None, :, :] @ p_local_exp[..., None]).squeeze(-1)
-            p_global += ecef2global_t_cast
-
-        return p_global
+        return self.global2local.apply_vectors_permutation(p_local)
 
     def _compute_local_from_global(self):
         if not self.uses_local_coords:
@@ -523,7 +717,7 @@ def _compute_local_from_global(self):
 
         p_global = np.stack([self.x_global, self.y_global, self.z_global], axis=-1)
 
-        p_local = self._vec_global2local(p_global)
+        p_local = self._global2local_points(p_global)
 
         # Store results
         self._x_local = p_local[:, 0]
@@ -535,8 +729,8 @@ def _compute_local_from_global(self):
         point_global = np.stack(polar_to_cartesian(
             r, self.pointn_azim_global, self.pointn_elev_global), axis=-1)
 
-        point_local_x, point_local_y, point_local_z = self._vec_global2local(
-            point_global, translate=False
+        point_local_x, point_local_y, point_local_z = self._global2local_vectors(
+            point_global
         ).T
 
         _, local_azimuth, local_elevation = cartesian_to_polar(
@@ -545,93 +739,23 @@ def _compute_local_from_global(self):
         self._pointn_elev_local = local_elevation
         self._pointn_azim_local = local_azimuth
 
-    def _vec_global2local(
-        self,
-        p_global,
-        *,
-        translate=True,
-        permutate=False
-    ):
-        """Receives a vector shaped as (N, 3) and applies global2local transform,
-            returning (N, 3).
-        If permutate == True, then may receive a vector (M, 3),
-            returning (N, M, 3)
-        """
-        N = self.num_geometries
-
-        global2ecef_t = self._global2ecef_transl_mtx if translate else 0
-        ecef2local_t = self._ecef2local_transl_mtx if translate else 0
-
-        if not permutate:
-            assert p_global.shape == (N, 3)
-
-            # remove ecef2global translation
-            p_global = p_global + global2ecef_t  # (N,3)
-
-            # rotate global2ecef2local
-            p_local = (self._global2local_rot_mtx @ p_global[..., None]).squeeze(-1)  # (N,3)
-
-            # add ecef2local translation
-            p_local += ecef2local_t  # (N,3)
-        else:
-            global2ecef_t_cast = global2ecef_t
-            ecef2local_t_cast = ecef2local_t
-            if translate:
-                global2ecef_t_cast = global2ecef_t[:, None, :]
-                ecef2local_t_cast = ecef2local_t[:, None, :]
-            p_global_exp = p_global[None, :, :] + global2ecef_t_cast  # (N,M,3)
-            p_local = (self._global2local_rot_mtx[:, None, :, :] @ p_global_exp[..., None]).squeeze(-1)
-            p_local += ecef2local_t_cast
-
-        return p_local
-
     def _compute_global_local_transform(self):
         # get ecef to local
-        local_lat, local_lon, local_alt = self.local_lla_references
-        rotation_around_z = -local_lon - 90
-        rotation_around_x = local_lat - 90
-
-        ecef2local_rot = scipy.spatial.transform.Rotation.from_euler(
-            'zx',
-            np.stack([rotation_around_z, rotation_around_x], axis=-1),
-            degrees=True
+        self.global2local = (
+            self.global_reference_frame
+                .to_ecef
+                .and_then(
+                    self.local_reference_frame.from_ecef
+                )
         )
-
-        # first translation, after rotation
-        self._local2ecef_transl_mtx = np.zeros((self.num_geometries, 3))
-        # NOTE: works because of spherical Earth
-        self._local2ecef_transl_mtx[:, 2] = local_alt + EARTH_RADIUS_M
-        local2ecef_rot_mtx = ecef2local_rot.inv().as_matrix()
-
-        # first rotation, after translation
-        ecef2local_rot_mtx = ecef2local_rot.as_matrix()
-        self._ecef2local_transl_mtx = -self._local2ecef_transl_mtx
-
-        # global transforms
-        global_lat, global_lon, global_alt = self._global_lla_reference
-        rotation_around_z = -global_lon - 90
-        rotation_around_x = global_lat - 90
-
-        ecef2global_rot = scipy.spatial.transform.Rotation.from_euler(
-            'zx',
-            np.stack([rotation_around_z, rotation_around_x], axis=-1),
-            degrees=True
+        self.local2global = (
+            self.local_reference_frame
+                .to_ecef
+                .and_then(
+                    self.global_reference_frame.from_ecef
+                )
         )
 
-        # first translation, after rotation
-        self._global2ecef_transl_mtx = np.zeros((1, 3))
-        # NOTE: works because of spherical Earth
-        self._global2ecef_transl_mtx[:, 2] = global_alt + EARTH_RADIUS_M
-        global2ecef_rot_mtx = ecef2global_rot.inv().as_matrix()
-
-        # first rotation, after translation
-        ecef2global_rot_mtx = ecef2global_rot.as_matrix()
-        self._ecef2global_transl_mtx = -self._global2ecef_transl_mtx
-
-        self._local2global_rot_mtx = ecef2global_rot_mtx @ local2ecef_rot_mtx
-
-        self._global2local_rot_mtx = ecef2local_rot_mtx @ global2ecef_rot_mtx
-
     def get_local_distance_to(
         self,
         other: "SimulatorGeometry",
@@ -655,9 +779,8 @@ def get_local_distance_to(
             return self.get_global_distance_to(other)
 
         p_global = np.stack([other.x_global, other.y_global, other.z_global], axis=-1)
-        other_local = self._vec_global2local(
+        other_local = self._global2local_points_permutation(
             p_global,
-            permutate=True
         )
         own_local = np.stack([self.x_local, self.y_local, np.zeros_like(self.z_local)], axis=-1)
 
@@ -765,17 +888,21 @@ def plot_geom(
 
 if __name__ == "__main__":
     global_lla = (-14, -45, 1200)
+    global_reference_frame = ENUReferenceFrame(
+        lat=np.array([global_lla[0]]),
+        lon=np.array([global_lla[1]]),
+        alt=np.array([global_lla[2]]),
+    )
     # tg = SimulatorGeometry(
     #     1,
     #     1,
-    #     global_lla
+    #     global_reference_frame
     # )
-    # tg.set_local_coord_sys(
+    # tg.set_local_reference_frame(
     #     np.array([-14]),
     #     np.array([-45]),
     #     np.array([1200]),
     # )
-    # print(tg.local_lla_references)
 
     from sharc.topology.topology_macrocell import TopologyMacrocell
     rng = np.random.RandomState(seed=0xcaffe)
@@ -790,23 +917,27 @@ def plot_geom(
     bs_geom = SimulatorGeometry(
         topology.num_base_stations,
         topology.num_base_stations,
-        global_lla
+        global_reference_frame
     )
     ue_geom = SimulatorGeometry(
         num_ue * topology.num_base_stations,
         num_ue * topology.num_base_stations,
-        global_lla
+        global_reference_frame
     )
 
-    bs_geom.set_local_coord_sys(
-        np.repeat(11, topology.num_base_stations),
-        np.repeat(-47, topology.num_base_stations),
-        np.repeat(1200, topology.num_base_stations),
+    bs_geom.set_local_reference_frame(
+        ENUReferenceFrame(
+            lat=np.repeat(11, topology.num_base_stations),
+            lon=np.repeat(-47, topology.num_base_stations),
+            alt=np.repeat(1200, topology.num_base_stations),
+        )
     )
-    ue_geom.set_local_coord_sys(
-        np.repeat(11, num_ue * topology.num_base_stations),
-        np.repeat(-47, num_ue * topology.num_base_stations),
-        np.repeat(1200, num_ue * topology.num_base_stations),
+    ue_geom.set_local_reference_frame(
+        ENUReferenceFrame(
+            lat=np.repeat(11, num_ue * topology.num_base_stations),
+            lon=np.repeat(-47, num_ue * topology.num_base_stations),
+            alt=np.repeat(1200, num_ue * topology.num_base_stations),
+        )
     )
 
     bs_geom.set_local_coords(
diff --git a/sharc/topology/topology_spherical_sampling_from_grid.py b/sharc/topology/topology_spherical_sampling_from_grid.py
index c96f4854..e9316295 100644
--- a/sharc/topology/topology_spherical_sampling_from_grid.py
+++ b/sharc/topology/topology_spherical_sampling_from_grid.py
@@ -1,7 +1,7 @@
 from sharc.topology.topology import Topology
 import numpy as np
 from sharc.support.sharc_geom import CoordinateSystem
-from sharc.support.geometry import SimulatorGeometry
+from sharc.support.geometry import SimulatorGeometry, ENUReferenceFrame
 from sharc.parameters.imt.parameters_grid import ParametersTerrestrialGrid
 # from sharc.satellite.utils.sat_utils import lla2ecef
 # import math
@@ -25,7 +25,7 @@ def __init__(
         max_ue_distance: float,
         num_base_stations: int,
         global_sim_lla_reference: tuple[float, float, float],
-        grid: ParametersTerrestrialGrid,
+        grid: ParametersTerrestrialGrid | np.ndarray,
     ):
         """
         Initializes a spherical topology with specific network settings.
@@ -61,13 +61,16 @@ def calculate_coordinates(self, random_number_gen=np.random.RandomState()):
         random_number_gen : np.random.RandomState, optional
             Random number generator (not used in this implementation).
         """
-        self.grid.reset_grid(
-            "calculate_coords",
-            random_number_gen,
-            True
-        )
+        if not isinstance(self.grid, np.ndarray):
+            self.grid.reset_grid(
+                "calculate_coords",
+                random_number_gen,
+                True
+            )
+            lla_grid_to_sample = self.grid.lon_lat_grid[::-1]
+        else:
+            lla_grid_to_sample = self.grid
 
-        lla_grid_to_sample = self.grid.lon_lat_grid[::-1]
         # print("self.grid.lon_lat_grid.shape", self.grid.lon_lat_grid.shape)
         # print("self.grid.lon_lat_grid.shape", self.grid.lon_lat_grid.shape)
         # print("self.num_base_stations", self.num_base_stations)
@@ -88,15 +91,21 @@ def calculate_coordinates(self, random_number_gen=np.random.RandomState()):
         geom = SimulatorGeometry(
             self.num_base_stations,
             self.num_base_stations,
-            self.global_cs,
+            ENUReferenceFrame(
+                lat=self.global_cs[0],
+                lon=self.global_cs[1],
+                alt=self.global_cs[2],
+            ),
         )
         lat, lon, alt = chosen_llas
         # coords locais para determinar
         # transformação global <> local
-        geom.set_local_coord_sys(
-            lat,
-            lon,
-            alt,
+        geom.set_local_reference_frame(
+            ENUReferenceFrame(
+                lat=lat,
+                lon=lon,
+                alt=alt,
+            )
         )
         self.x = np.zeros(self.num_base_stations)
         self.y = np.zeros(self.num_base_stations)
@@ -128,10 +137,12 @@ def get_ue_geometry(self, ue_k: int) -> SimulatorGeometry:
             self.num_base_stations * ue_k,
             self.global_cs,
         )
-        ue_geom.set_local_coord_sys(
-            np.repeat(self.bs_geometry._local_lla_references[0], ue_k),
-            np.repeat(self.bs_geometry._local_lla_references[1], ue_k),
-            np.repeat(self.bs_geometry._local_lla_references[2], ue_k),
+        ue_geom.set_local_reference_frame(
+            ENUReferenceFrame(
+                lat=np.repeat(self.bs_geometry._local_lla_references[0], ue_k),
+                lon=np.repeat(self.bs_geometry._local_lla_references[1], ue_k),
+                alt=np.repeat(self.bs_geometry._local_lla_references[2], ue_k),
+            )
         )
         return ue_geom
 
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index 5d8a1f32..a9f68c11 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -1,9 +1,341 @@
 import unittest
 import numpy as np
 import numpy.testing as npt
-from sharc.support.geometry import SimulatorGeometry
+from sharc.support.geometry import (
+    SimulatorGeometry, ENUReferenceFrame, RigidTransform
+)
 from sharc.satellite.ngso.constants import EARTH_RADIUS_M
 from copy import deepcopy
+from scipy.spatial.transform import Rotation
+from itertools import product
+
+
+def random_rigid_transform(rng, N):
+    """
+    Generate a random RigidTransform with batch size N.
+    Rotations are orthonormal matrices generated via QR.
+    """
+    A = rng.normal(size=(N, 3, 3))
+    Q, _ = np.linalg.qr(A)
+
+    # ensure right-handed (det = +1)
+    det = np.linalg.det(Q)
+    Q[det < 0, :, 0] *= -1
+
+    t = rng.normal(size=(N, 3))
+
+    return RigidTransform(Rotation.from_matrix(Q), t)
+
+def rot_identity(n=1):
+    return Rotation.from_rotvec(np.zeros((n, 3)))
+
+def rot_z(angle_deg, n=1):
+    return Rotation.from_rotvec(
+        np.tile([0.0, 0.0, angle_deg], (n, 1)),
+        degrees=True
+    )
+
+def rot_x(angle_deg, n=1):
+    return Rotation.from_rotvec(
+        np.tile([angle_deg, 0.0, 0.0], (n, 1)),
+        degrees=True
+    )
+
+def rot_y(angle_deg, n=1):
+    return Rotation.from_rotvec(
+        np.tile([0.0, angle_deg, 0.0], (n, 1)),
+        degrees=True
+    )
+
+
+class TestRigidTransform(unittest.TestCase):
+    def setUp(self):
+        pass
+
+    def test_init_and_broadcasting(self):
+        for rot_shp, t_shp in product(
+            [(1, 3), (4, 3)],
+            [(1, 3), (4, 3)],
+        ):
+            Nrot = rot_shp[0]
+            Nt = t_shp[0]
+            N = max(Nrot, Nt)
+
+            rot = rot_identity(rot_shp[0])
+            t = np.zeros(t_shp)
+
+            with self.assertRaises(ValueError):
+                RigidTransform(rot, np.zeros((3,)))
+
+            if Nt > 1:
+                with self.assertRaises(ValueError):
+                    RigidTransform(rot_identity(N+1), t)
+
+            if Nrot > 1:
+                with self.assertRaises(ValueError):
+                    RigidTransform(rot, np.zeros((N+1, 3)))
+
+            # should not throw:
+            tr = RigidTransform(rot, t)
+            RigidTransform(rot, np.zeros((1, 3)))
+            RigidTransform(rot_identity(1), t)
+
+            # test broadcasting contracts:
+            for fn in [tr.apply_points, tr.apply_vectors]:
+                if N > 1:
+                    with self.assertRaises(AssertionError):
+                        fn(np.zeros((N+1, 3)))
+
+                for in_shp in [(1, 3), (N, 3)]:
+                    res = fn(np.zeros(in_shp))
+                    npt.assert_equal(
+                        res.shape, (N, 3)
+                    )
+
+            for fn, in_shp in product(
+                [tr.apply_points_permutation, tr.apply_vectors_permutation],
+                [(1, 3), (N, 3), (N+1, 3)]
+            ):
+                Nin = in_shp[0]
+                res = fn(np.zeros(in_shp))
+                npt.assert_equal(
+                    res.shape, (N, Nin, 3)
+                )
+
+    def test_simple_transformations(self):
+        eps = 1e-4
+
+        ux = np.array([1., 0., 0.])
+        uy = np.array([0., 1., 0.])
+        uz = np.array([0., 0., 1.])
+        u = np.array([ux, uy, uz, ux + uy + uz])
+
+        #######################################################################
+        # Identity transform
+        id_tr = RigidTransform(rot_identity(4), np.zeros((1, 3)))
+        out = id_tr.apply_points(u)
+        npt.assert_equal(out, u)
+
+        out = id_tr.apply_vectors(u)
+        npt.assert_equal(out, u)
+
+        out = id_tr.apply_points_permutation(u)
+        npt.assert_equal(out, np.stack((u, u, u, u)))
+
+        out = id_tr.apply_vectors_permutation(u)
+        npt.assert_equal(out, np.stack((u, u, u, u)))
+
+        #######################################################################
+        # Translation only transform
+        t1_tr = RigidTransform(rot_identity(4), np.ones((1, 3)))
+        out = t1_tr.apply_points(u)
+        npt.assert_equal(out, u + 1.)
+        out = t1_tr.inv().apply_points(out)
+        npt.assert_equal(out, u)
+
+        out = t1_tr.apply_vectors(u)
+        npt.assert_equal(out, u)
+
+        out = t1_tr.apply_points_permutation(u)
+        npt.assert_equal(out, np.stack((u, u, u, u)) + 1.)
+        out = t1_tr.inv().apply_points(out[0])
+        npt.assert_equal(out, u)
+
+        out = t1_tr.apply_vectors_permutation(u)
+        npt.assert_equal(out, np.stack((u, u, u, u)))
+
+        #######################################################################
+        # Rotation Z only transform
+        rot_z90_tr = RigidTransform(
+            rot_z(-90.), np.zeros((1, 3))
+        )
+        out_rot_z90 = rot_z90_tr.apply_points(u)
+        npt.assert_allclose(
+            out_rot_z90,
+            np.array([
+                -uy, ux, uz, ux - uy + uz
+            ]),
+            atol=eps
+        )
+        out = rot_z90_tr.inv().apply_points(out_rot_z90)
+        npt.assert_allclose(
+            out,
+            u,
+            atol=eps
+        )
+
+        #######################################################################
+        # Rotation X only transform
+        rot_x90_tr = RigidTransform(
+            rot_x(-90., 4), np.zeros((1, 3))
+        )
+
+        out_rot_z90_x90 = rot_x90_tr.apply_points(out_rot_z90)
+        expected_out_rot_z90_x90 = np.array([
+            uz, ux, uy, ux + uy + uz
+        ])
+        npt.assert_allclose(
+            out_rot_z90_x90,
+            expected_out_rot_z90_x90,
+            atol=eps
+        )
+        out = rot_x90_tr.inv().apply_points(expected_out_rot_z90_x90)
+        npt.assert_allclose(
+            out,
+            out_rot_z90,
+            atol=eps
+        )
+
+        #######################################################################
+        # Rotation Z and then Rotation X transform
+        npt.assert_allclose(
+            rot_z90_tr.and_then(rot_x90_tr).apply_points(u),
+            expected_out_rot_z90_x90,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (rot_z90_tr.and_then(rot_x90_tr).inv()
+                .apply_points(expected_out_rot_z90_x90)),
+            u,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (rot_x90_tr.inv().and_then(rot_z90_tr.inv())
+                .apply_points(expected_out_rot_z90_x90)),
+            u,
+            atol=eps,
+        )
+
+        #######################################################################
+        # Rotation Z and then Rotation X and then Translation transform
+        npt.assert_allclose(
+            rot_z90_tr.and_then(rot_x90_tr).and_then(t1_tr).apply_points(u),
+            expected_out_rot_z90_x90 + 1.,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (rot_z90_tr.and_then(rot_x90_tr).and_then(t1_tr).inv()
+                .apply_points(expected_out_rot_z90_x90 + 1.)),
+            u,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (t1_tr.inv().and_then(rot_x90_tr.inv().and_then(rot_z90_tr.inv()))
+                .apply_points(expected_out_rot_z90_x90 + 1.)),
+            u,
+            atol=eps,
+        )
+        # considering that VECTOR calculation should NOT translate
+        npt.assert_allclose(
+            rot_z90_tr.and_then(rot_x90_tr).and_then(t1_tr).apply_vectors(u),
+            expected_out_rot_z90_x90,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (rot_z90_tr.and_then(rot_x90_tr).and_then(t1_tr).inv()
+                .apply_vectors(expected_out_rot_z90_x90)),
+            u,
+            atol=eps,
+        )
+        npt.assert_allclose(
+            (t1_tr.inv().and_then(rot_x90_tr.inv().and_then(rot_z90_tr.inv()))
+                .apply_vectors(expected_out_rot_z90_x90)),
+            u,
+            atol=eps,
+        )
+
+    def test_permutation_points_equivalence(self):
+        rng = np.random.default_rng(0)
+
+        for n in range(1, 10):
+            tr = random_rigid_transform(rng, n)
+
+            x = rng.normal(size=(n, 3))
+
+            y = tr.apply_points(x)
+            y_permutation = tr.apply_points_permutation(x)
+
+            # diagonal of permutation must equal non-permutation
+            npt.assert_allclose(
+                y,
+                y_permutation.diagonal().T,
+                atol=1e-12,
+            )
+
+    def test_permutation_vectors_equivalence(self):
+        rng = np.random.default_rng(0)
+
+        for n in range(1, 10):
+            tr = random_rigid_transform(rng, n)
+
+            x = rng.normal(size=(n, 3))
+
+            y = tr.apply_vectors(x)
+            y_permutation = tr.apply_vectors_permutation(x)
+
+            # diagonal of permutation must equal non-permutation
+            npt.assert_allclose(
+                y,
+                y_permutation.diagonal().T,
+                atol=1e-12,
+            )
+
+    def test_take_commutes_with_apply_points(self):
+        rng = np.random.default_rng(0)
+
+        for n in range(1, 10):
+            tr = random_rigid_transform(rng, n)
+            x = rng.normal(size=(n, 3))
+
+            y_full = tr.apply_points(x)
+
+            for i in range(n):
+                y_take = tr.take(i).apply_points(x)
+                npt.assert_allclose(
+                    y_take[i],
+                    y_full[i],
+                    atol=1e-12,
+                )
+
+    def test_take_matches_permutation_points(self):
+        rng = np.random.default_rng(1)
+
+        for n in [2, 5]:
+            tr = random_rigid_transform(rng, n)
+            x = rng.normal(size=(n, 3))
+
+            Y = tr.apply_points_permutation(x)
+
+            for i in range(n):
+                tr_i = tr.take(i)
+                yi = tr_i.apply_points(x)
+
+                npt.assert_allclose(
+                    yi,
+                    Y[i],
+                    rtol=1e-12,
+                    atol=1e-12,
+                )
+
+    def test_take_matches_permutation_vectors(self):
+        rng = np.random.default_rng(1)
+
+        for n in [2, 5]:
+            tr = random_rigid_transform(rng, n)
+            x = rng.normal(size=(n, 3))
+
+            Y = tr.apply_vectors_permutation(x)
+
+            for i in range(n):
+                tr_i = tr.take(i)
+                yi = tr_i.apply_vectors(x)
+
+                npt.assert_allclose(
+                    yi,
+                    Y[i],
+                    rtol=1e-12,
+                    atol=1e-12,
+                )
 
 
 class TestGeometry(unittest.TestCase):
@@ -18,9 +350,14 @@ def _test_expected_geom(
         only_global_azim=None,
         only_local_azim=None
     ):
-        npt.assert_allclose(geom.x_local, expect_local["x"], atol=0.001)
-        npt.assert_allclose(geom.y_local, expect_local["y"], atol=0.001)
-        npt.assert_allclose(geom.z_local, expect_local["z"], atol=0.001)
+        npt.assert_allclose(
+            np.stack((geom.x_local, geom.y_local, geom.z_local)),
+            np.stack((expect_local["x"], expect_local["y"], expect_local["z"])),
+            atol=0.001
+        )
+        # npt.assert_allclose(geom.x_local, expect_local["x"], atol=0.001)
+        # npt.assert_allclose(geom.y_local, expect_local["y"], atol=0.001)
+        # npt.assert_allclose(geom.z_local, expect_local["z"], atol=0.001)
         npt.assert_allclose(geom.pointn_elev_local, expect_local["elev"], atol=0.001)
         if only_local_azim is None:
             npt.assert_allclose(geom.pointn_azim_local, expect_local["azim"], atol=0.001)
@@ -32,9 +369,14 @@ def _test_expected_geom(
                 atol=0.001
             )
 
-        npt.assert_allclose(geom.x_global, expect_global["x"], atol=0.001)
-        npt.assert_allclose(geom.y_global, expect_global["y"], atol=0.001)
-        npt.assert_allclose(geom.z_global, expect_global["z"], atol=0.001)
+        npt.assert_allclose(
+            np.stack((geom.x_global, geom.y_global, geom.z_global)),
+            np.stack((expect_global["x"], expect_global["y"], expect_global["z"])),
+            atol=0.001
+        )
+        # npt.assert_allclose(geom.x_global, expect_global["x"], atol=0.001)
+        # npt.assert_allclose(geom.y_global, expect_global["y"], atol=0.001)
+        # npt.assert_allclose(geom.z_global, expect_global["z"], atol=0.001)
         npt.assert_allclose(geom.pointn_elev_global, expect_global["elev"], atol=0.001)
         if only_global_azim is None:
             npt.assert_allclose(
@@ -57,14 +399,18 @@ def setUp(self):
     def test_set_coords_when_local_eq_global(self):
         """Test setting coordinates when local should eq global
         """
-        ref = (10, -5, 1200)
-
-        no_local = SimulatorGeometry(3, False, ref)
-        local_eq_global = SimulatorGeometry(3, True, ref)
-        local_eq_global.set_local_coord_sys(
-            np.repeat([ref[0]], 3),
-            np.repeat([ref[1]], 3),
-            np.repeat([ref[2]], 3),
+        ref_lla = (10, -5, 1200)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
+        no_local = SimulatorGeometry(3, False, ref_frame)
+        local_eq_global = SimulatorGeometry(3, True, ref_frame)
+        local_eq_global.set_local_reference_frame(
+            ENUReferenceFrame(
+                lat=np.repeat([ref_lla[0]], 3),
+                lon=np.repeat([ref_lla[1]], 3),
+                alt=np.repeat([ref_lla[2]], 3),
+            )
         )
 
         """Test setting global coordinates
@@ -102,16 +448,21 @@ def test_set_coords_when_local_eq_global(self):
     def test_setting_different_alts(self):
         """Test setting coordinates when local should eq global
         """
-        ref = (-45.0, 90, 30.)
+        ref_lla = (-45.0, 90, 30.)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
 
         alt_vals = np.array([500., 800., 1200., 1800.])
 
         def init_diff_alt():
-            geom = SimulatorGeometry(4, True, ref)
-            geom.set_local_coord_sys(
-                np.repeat([ref[0]], 4),
-                np.repeat([ref[1]], 4),
-                alt_vals,
+            geom = SimulatorGeometry(4, True, ref_frame)
+            geom.set_local_reference_frame(
+                ENUReferenceFrame(
+                    lat=np.repeat([ref_lla[0]], 4),
+                    lon=np.repeat([ref_lla[1]], 4),
+                    alt=alt_vals,
+                )
             )
             return geom
         diff_alt = init_diff_alt()
@@ -124,7 +475,7 @@ def init_diff_alt():
             "azim": np.array([90., 10., -179., 180.]), "elev": np.array([89., -89., -1., 12.]),
         }
         expect_global = deepcopy(expect_local)
-        alt_vals_diff = ref[2] - alt_vals
+        alt_vals_diff = ref_lla[2] - alt_vals
         expect_global["z"] -= alt_vals_diff
 
         diff_alt.set_global_coords(**expect_global)
@@ -147,7 +498,10 @@ def init_diff_alt():
     def test_setting_different_llas(self):
         """Test setting coordinates when local should eq global
         """
-        ref = (0.0, 0.0, 0.0)
+        ref_lla = (0.0, 0.0, 0.0)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
 
         llas = np.array([
             [-90., 0., 0.],
@@ -157,11 +511,13 @@ def test_setting_different_llas(self):
         ]).T
 
         def init_diff_llas():
-            geom = SimulatorGeometry(4, True, ref)
-            geom.set_local_coord_sys(
-                llas[0],
-                llas[1],
-                llas[2],
+            geom = SimulatorGeometry(4, True, ref_frame)
+            geom.set_local_reference_frame(
+                ENUReferenceFrame(
+                    lat=llas[0],
+                    lon=llas[1],
+                    alt=llas[2],
+                )
             )
             return geom
 
@@ -203,7 +559,10 @@ def test_get_local_distance_to_diff_ref(self):
         """Tests getting local distance from a station to another when
         they have different local references
         """
-        ref = (90., 0., 0.)
+        ref_lla = (90., 0., 0.)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
         local_llas = np.array([
             [0., 0., 0.],
             [0., 1., 0.],
@@ -213,11 +572,13 @@ def test_get_local_distance_to_diff_ref(self):
         ]).T
 
         def init_geom():
-            geom = SimulatorGeometry(5, True, ref)
-            geom.set_local_coord_sys(
-                local_llas[0],
-                local_llas[1],
-                local_llas[2],
+            geom = SimulatorGeometry(5, True, ref_frame)
+            geom.set_local_reference_frame(
+                ENUReferenceFrame(
+                    lat=local_llas[0],
+                    lon=local_llas[1],
+                    alt=local_llas[2],
+                )
             )
             geom.set_local_coords(
                 np.repeat(0., 5),
@@ -268,7 +629,10 @@ def test_get_local_distance_to_same_ref(self):
         """Tests getting local distance from a station to another when
         they have same local references
         """
-        ref = (90., 0., 0.)
+        ref_lla = (90., 0., 0.)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
         local_llas = np.array([
             [0., 0., 0.],
             [0., 0., 0.],
@@ -277,11 +641,13 @@ def test_get_local_distance_to_same_ref(self):
         ]).T
 
         def init_geom():
-            geom = SimulatorGeometry(4, True, ref)
-            geom.set_local_coord_sys(
-                local_llas[0],
-                local_llas[1],
-                local_llas[2],
+            geom = SimulatorGeometry(4, True, ref_frame)
+            geom.set_local_reference_frame(
+                ENUReferenceFrame(
+                    lat=local_llas[0],
+                    lon=local_llas[1],
+                    alt=local_llas[2],
+                )
             )
             geom.set_local_coords(
                 np.tile([0., 20.], 2),
@@ -353,7 +719,10 @@ def init_geom():
     def test_get_local_elevation(self):
         """Tests getting local elevation from a station to another.
         """
-        ref = (90., 0., 0.)
+        ref_lla = (90., 0., 0.)
+        ref_frame = ENUReferenceFrame(
+            lat=ref_lla[0], lon=ref_lla[1], alt=ref_lla[2],
+        )
         local_llas = np.array([
             [0., 0., 0.],
             [0., 0., 0.],
@@ -362,11 +731,13 @@ def test_get_local_elevation(self):
         ]).T
 
         def init_geom():
-            geom = SimulatorGeometry(4, True, ref)
-            geom.set_local_coord_sys(
-                local_llas[0],
-                local_llas[1],
-                local_llas[2],
+            geom = SimulatorGeometry(4, True, ref_frame)
+            geom.set_local_reference_frame(
+                ENUReferenceFrame(
+                    lat=local_llas[0],
+                    lon=local_llas[1],
+                    alt=local_llas[2],
+                )
             )
             geom.set_local_coords(
                 np.tile([0., 20.], 2),

From 88902e29b3a47d15dd2900caf2756324a6feed1d Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 16:52:31 -0300
Subject: [PATCH 03/10] hotfix: topology sampling from spherical grid

---
 sharc/topology/topology_spherical_sampling_from_grid.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/sharc/topology/topology_spherical_sampling_from_grid.py b/sharc/topology/topology_spherical_sampling_from_grid.py
index e9316295..552b0f43 100644
--- a/sharc/topology/topology_spherical_sampling_from_grid.py
+++ b/sharc/topology/topology_spherical_sampling_from_grid.py
@@ -98,6 +98,7 @@ def calculate_coordinates(self, random_number_gen=np.random.RandomState()):
             ),
         )
         lat, lon, alt = chosen_llas
+        self.chosen_lat, self.chosen_lon, self.chosen_alt = lat, lon, alt
         # coords locais para determinar
         # transformação global <> local
         geom.set_local_reference_frame(
@@ -135,13 +136,13 @@ def get_ue_geometry(self, ue_k: int) -> SimulatorGeometry:
         ue_geom = SimulatorGeometry(
             self.num_base_stations * ue_k,
             self.num_base_stations * ue_k,
-            self.global_cs,
+            self.bs_geometry.global_reference_frame,
         )
         ue_geom.set_local_reference_frame(
             ENUReferenceFrame(
-                lat=np.repeat(self.bs_geometry._local_lla_references[0], ue_k),
-                lon=np.repeat(self.bs_geometry._local_lla_references[1], ue_k),
-                alt=np.repeat(self.bs_geometry._local_lla_references[2], ue_k),
+                lat=np.repeat(self.chosen_lat, ue_k),
+                lon=np.repeat(self.chosen_lon, ue_k),
+                alt=np.repeat(self.chosen_alt, ue_k),
             )
         )
         return ue_geom

From 8df084d601b0c0b91b68be833309c916ee9f8cac Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 16:54:07 -0300
Subject: [PATCH 04/10] update: array antenna for local coordinate system AND
 faster computation

---
 sharc/antenna/antenna_array.py | 145 ++++++++++++++++++++++-----------
 1 file changed, 99 insertions(+), 46 deletions(-)

diff --git a/sharc/antenna/antenna_array.py b/sharc/antenna/antenna_array.py
index 582633be..18dea309 100644
--- a/sharc/antenna/antenna_array.py
+++ b/sharc/antenna/antenna_array.py
@@ -9,6 +9,7 @@
 from sharc.antenna.antenna import Antenna
 from sharc.parameters.imt.parameters_antenna_imt import ParametersAntennaImt
 from sharc.support.geometry import RigidTransform
+from sharc.support.sharc_geom import polar_to_cartesian, cartesian_to_polar
 
 import numpy as np
 
@@ -19,10 +20,26 @@ class AntennaArray(Antenna):
     def __init__(
         self,
         par: ParametersAntennaImt,
-        # global2local_transform: RigidTransform
+        global2local_transform: RigidTransform = None,
     ):
         super().__init__()
         self.par = par
+        self.always_first_beam = False
+
+        self.global2local_transform = global2local_transform
+        if self.global2local_transform is not None:
+            if self.global2local_transform.N > 1:
+                raise ValueError(
+                    "global2local_transform is supposed to have a single"
+                    " transformation for the purposes of antenna calculations"
+                )
+
+    def set_always_first_beam(self):
+        """
+        In case this is called, then calculate_gains will sum all beams
+        contributions for each direction angle.
+        """
+        self.always_first_beam = True
 
     def calculate_gain(self, *args, **kwargs) -> np.array:
         """
@@ -35,7 +52,9 @@ def calculate_gain(self, *args, **kwargs) -> np.array:
             self.par.adjacent_antenna_model == "SINGLE_ELEMENT"
             and not co_channel
         )
-        if "beams_l" in kwargs.keys():
+        if self.always_first_beam:
+            beam_idxs = np.zeros(len(phi_vec), dtype=int)
+        elif "beams_l" in kwargs.keys():
             beam_idxs = np.asarray(kwargs["beams_l"], dtype=int)
         else:
             beam_idxs = np.arange(len(phi_vec))
@@ -74,29 +93,35 @@ def _array_gain(
         phi: np.ndarray, theta: np.ndarray,
         beam_idxs: np.ndarray
     ):
-        v_vec = self._super_position_vector(
-            phi, theta,
-            self.par.n_rows, self.par.n_columns,
-            self.par.element_vert_spacing,
-            self.par.element_horiz_spacing,
-        )
         if len(self.beams_list) == 0:
             beam_phi, beam_theta = phi, theta
         else:
             beam_phi, beam_theta = np.array(self.beams_list).T
+            beam_phi, beam_theta = self._to_local_coord(beam_phi, beam_theta)
 
         beam_etilt = beam_theta - 90.
-        beams_w_vec = self._weight_vector(
+        beams_w_vec_row, beams_w_vec_col = self._weight_vector(
             beam_phi, beam_etilt,
             self.par.n_rows, self.par.n_columns,
             self.par.element_vert_spacing,
             self.par.element_horiz_spacing,
         )
-        w_vec = beams_w_vec[beam_idxs]
+        w_vec_row, w_vec_col = beams_w_vec_row[beam_idxs], beams_w_vec_col[beam_idxs]
+
+        v_vec_row, v_vec_col = self._super_position_vector(
+            phi, theta,
+            self.par.n_rows, self.par.n_columns,
+            self.par.element_vert_spacing,
+            self.par.element_horiz_spacing,
+        )
 
+        # NOTE: this formula has the same result to the one presented on M.2101
+        # but it is optimized for computation
+        # considering W(m, n) = W(m)W(n) and V(m, n) = V(m)V(n)
         g = 10 * np.log10(
             abs(
-                np.sum(v_vec * w_vec, axis=(1, 2))
+                np.sum(v_vec_row * w_vec_row, axis=-1)
+                * np.sum(v_vec_col * w_vec_col, axis=-1)
             )**2
         )
 
@@ -120,29 +145,33 @@ def _weight_vector(
 
         Returns
         -------
-            w_vec (np.array): weighting vector
+            w_vec (np.array, np.array):
+                weighting vectors, first for rows, second for columns
         """
         # shape (Na, 1, 1)
         r_phi = np.atleast_1d(
             np.deg2rad(phi_tilt)
-        )[:, np.newaxis, np.newaxis]
+        )[:, np.newaxis]
         r_theta = np.atleast_1d(
             np.deg2rad(theta_tilt)
-        )[:, np.newaxis, np.newaxis]
+        )[:, np.newaxis]
 
         # shape (1, Nr, 1)
-        n = np.arange(n_rows)[np.newaxis, :, np.newaxis] + 1
+        n = np.arange(n_rows)[np.newaxis, :] + 1
         # shape (1, 1, Nc)
-        m = np.arange(n_cols)[np.newaxis, np.newaxis, :] + 1
+        m = np.arange(n_cols)[np.newaxis, :] + 1
+
+        exp_arg_n = (n - 1) * dv * np.sin(r_theta)
+        exp_arg_m = (m - 1) * dh * np.cos(r_theta) * np.sin(r_phi)
 
-        exp_arg = (n - 1) * dv * np.sin(r_theta) - \
-                  (m - 1) * dh * np.cos(r_theta) * np.sin(r_phi)
+        w_vec_n = (1 / np.sqrt(n_rows * n_cols)) *\
+            np.exp(2 * np.pi * 1.0j * exp_arg_n)
 
-        w_vec = (1 / np.sqrt(n_rows * n_cols)) *\
-            np.exp(2 * np.pi * 1.0j * exp_arg)
+        w_vec_m = (1 / np.sqrt(n_rows * n_cols)) *\
+            np.exp(2 * np.pi * 1.0j * exp_arg_m)
 
         # shape (Na, Nr, Nc)
-        return w_vec
+        return (w_vec_n, w_vec_m)
 
     @staticmethod
     def _super_position_vector(
@@ -163,26 +192,27 @@ def _super_position_vector(
         -------
             v_vec (np.array): superposition vector
         """
-        # shape (Na, 1, 1)
-        r_phi = np.atleast_1d(
-            np.deg2rad(phi)
-        )[:, np.newaxis, np.newaxis]
-        r_theta = np.atleast_1d(
-            np.deg2rad(theta)
-        )[:, np.newaxis, np.newaxis]
+        phi = np.atleast_1d(phi)
+        theta = np.atleast_1d(theta)
 
-        # shape (1, Nr, 1)
-        n = np.arange(n_rows)[np.newaxis, :, np.newaxis] + 1
-        # shape (1, 1, Nc)
-        m = np.arange(n_cols)[np.newaxis, np.newaxis, :] + 1
+        # (Na,)
+        A = dv * np.cos(np.deg2rad(theta))
+        B = dh * np.sin(np.deg2rad(theta)) * np.sin(np.deg2rad(phi))
 
-        exp_arg = (n - 1) * dv * np.cos(r_theta) + \
-                  (m - 1) * dh * np.sin(r_theta) * np.sin(r_phi)
+        # indices
+        n = np.arange(n_rows)          # (Nr,)
+        m = np.arange(n_cols)          # (Nc,)
 
-        v_vec = np.exp(2 * np.pi * 1.0j * exp_arg)
+        # (Na, Nr)
+        row_phase = np.exp(
+            2j * np.pi * A[:, None] * n[None, :]
+        )
+        # (Na, Nc)
+        col_phase = np.exp(
+            2j * np.pi * B[:, None] * m[None, :]
+        )
 
-        # shape (Na, Nr, Nc)
-        return v_vec
+        return (row_phase, col_phase)
 
     @staticmethod
     def _element_gain_dispatch(par: ParametersAntennaImt, phi, theta):
@@ -212,7 +242,7 @@ def _calculate_m2101_element_gain(
             multiplication_factor * (phi / phi_3db)**2, am
         )
         g_vertical = -1.0 * np.minimum(
-            multiplication_factor * ((theta-90.) / theta_3db)**2, sla_v
+            multiplication_factor * ((theta - 90.) / theta_3db)**2, sla_v
         )
 
         att = -1.0 * (
@@ -223,7 +253,20 @@ def _calculate_m2101_element_gain(
         return g_max - np.minimum(att, am)
 
     def _to_local_coord(self, phi, theta):
-        return np.array(phi), np.array(theta)
+        if self.global2local_transform is None:
+            return np.array(phi), np.array(theta)
+
+        theta_from_plane = 90 - theta
+        vecs = np.stack(polar_to_cartesian(1, phi, theta_from_plane), axis=-1)
+        transformed_vecs = self.global2local_transform.apply_vectors(
+            vecs
+        )
+        x, y, z = transformed_vecs.T
+        _r, phi, elev_from_plane = cartesian_to_polar(x, y, z)
+
+        theta = 90 - elev_from_plane
+
+        return phi, theta
 
     def add_beam(self, phi_etilt: float, theta_etilt: float):
         """
@@ -236,10 +279,11 @@ def add_beam(self, phi_etilt: float, theta_etilt: float):
             phi_etilt (float): azimuth electrical tilt angle [degrees]
             theta_etilt (float): elevation electrical tilt angle [degrees]
         """
-        # phi_etilt, theta_etilt = np.atleast_1d(phi_etilt), np.atleast_1d(theta_etilt)
-        phi, theta = self._to_local_coord(phi_etilt, theta_etilt)
+        phi_etilt, theta_etilt = np.atleast_1d(phi_etilt), np.atleast_1d(theta_etilt)
+
+    # def add_beam_in_local_coords(self):
         self.beams_list.append(
-            (np.ndarray.item(phi), np.ndarray.item(theta)),
+            (np.ndarray.item(phi_etilt), np.ndarray.item(theta_etilt)),
         )
 
 
@@ -262,6 +306,13 @@ def add_beam(self, phi_etilt: float, theta_etilt: float):
     antenna_params.multiplication_factor = 12
 
     par = antenna_params.get_antenna_parameters()
+    # from sharc.support.geometry import ENUReferenceFrame
+    # ref_frame = ENUReferenceFrame(
+    #     lat=np.array([90.]),
+    #     lon=np.array([-90.]),
+    #     alt=np.array([0.]),
+    # )
+    # antenna = AntennaArray(par, ref_frame.from_ecef)
     antenna = AntennaArray(par)
 
     antenna.add_beam(np.array(0.), np.array(90.))
@@ -285,9 +336,9 @@ def add_beam(self, phi_etilt: float, theta_etilt: float):
     ax1 = fig.add_subplot(121)
 
     ax1.plot(phi_scan, gain)
-    top_y_lim = np.ceil(np.max(gain) / 10) * 10
+    top_y_lim1 = np.ceil(np.max(gain) / 10) * 10
     ax1.set_xlim(-180, 180)
-    ax1.set_ylim(top_y_lim - 60, top_y_lim)
+    ax1.set_ylim(top_y_lim1 - 60, top_y_lim1)
     ax1.grid(True)
     ax1.set_xlabel(r"$\varphi$ [deg]")
     ax1.set_ylabel("Gain [dBi]")
@@ -308,9 +359,11 @@ def add_beam(self, phi_etilt: float, theta_etilt: float):
     ax2 = fig.add_subplot(122, sharey=ax1)
 
     ax2.plot(theta_scan, gain)
-    top_y_lim = np.ceil(np.max(gain) / 10) * 10
+    top_y_lim2 = np.ceil(np.max(gain) / 10) * 10
+    top_y_lim2 = np.maximum(top_y_lim1, top_y_lim2)
+
     ax2.set_xlim(0, 180.)
-    ax2.set_ylim(top_y_lim - 60, top_y_lim)
+    ax2.set_ylim(top_y_lim2 - 60, top_y_lim2)
     ax2.grid(True)
     ax2.set_xlabel(r"$\vartheta$ [deg]")
     ax2.set_ylabel("Gain [dBi]")

From 2622b8c2ee2a8efd24ebe668ce84020aefb984ff Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 16:58:05 -0300
Subject: [PATCH 05/10] update: antenna array2 in generate_mss_d2d &
 DWNReferenceFrame for it working

---
 sharc/parameters/parameters_antenna.py |  6 ++-
 sharc/station_factory.py               | 46 +++++++++++++++++--
 sharc/support/geometry.py              | 42 ++++++++++++++---
 tests/test_geometry.py                 | 62 +++++++++++++++++++++++---
 4 files changed, 141 insertions(+), 15 deletions(-)

diff --git a/sharc/parameters/parameters_antenna.py b/sharc/parameters/parameters_antenna.py
index a8fc8a0e..189514c9 100644
--- a/sharc/parameters/parameters_antenna.py
+++ b/sharc/parameters/parameters_antenna.py
@@ -26,6 +26,7 @@ class ParametersAntenna(ParametersBase):
         "ITU-R S.1855",
         "ITU-R Reg. RR. Appendice 7 Annex 3",
         "ARRAY",
+        "ARRAY2",
         "ITU-R-S.1528-Taylor",
         "ITU-R-S.1528-Section1.2",
         "ITU-R-S.1528-LEO",
@@ -41,6 +42,7 @@ class ParametersAntenna(ParametersBase):
                             "ITU-R S.1855",
                             "ITU-R Reg. RR. Appendice 7 Annex 3",
                             "ARRAY",
+                            "ARRAY2",
                             "ITU-R-S.1528-Taylor",
                             "ITU-R-S.1528-Section1.2",
                             "ITU-R-S.1528-LEO",
@@ -150,7 +152,7 @@ def validate(self, ctx):
                 f"{ctx}.pattern should be set. Is None instead",
             )
 
-        if self.pattern != "ARRAY" and self.gain is None:
+        if self.pattern != "ARRAY" and self.pattern != "ARRAY2" and self.gain is None:
             raise ValueError(
                 f"{ctx}.gain should be set if not using array antenna.",
             )
@@ -182,7 +184,7 @@ def validate(self, ctx):
                     # just hijacking validation since diameter is optional
                     self.itu_reg_rr_a7_3.diameter = 0
                 self.itu_reg_rr_a7_3.validate(f"{ctx}.itu_reg_rr_a7_3")
-            case "ARRAY":
+            case "ARRAY" | "ARRAY2":
                 # TODO: validate here and make array non imt specific
                 # self.array.validate(
                 #     f"{ctx}.array",
diff --git a/sharc/station_factory.py b/sharc/station_factory.py
index 02384d4f..a6666863 100644
--- a/sharc/station_factory.py
+++ b/sharc/station_factory.py
@@ -10,6 +10,9 @@
 import sys
 import math
 
+from sharc.support.geometry import (
+    DWNReferenceFrame, SimulatorGeometry, ENUReferenceFrame
+)
 from sharc.support.enumerations import StationType
 from sharc.parameters.parameters import Parameters
 from sharc.parameters.imt.parameters_imt import ParametersImt
@@ -44,6 +47,7 @@
 from sharc.antenna.antenna_rs1861_9c import AntennaRS1861_9C
 from sharc.antenna.antenna_rs2043 import AntennaRS2043
 from sharc.antenna.antenna_s465 import AntennaS465
+from sharc.antenna.antenna_array import AntennaArray
 from sharc.antenna.antenna_rra7_3 import AntennaReg_RR_A7_3
 from sharc.antenna.antenna_modified_s465 import AntennaModifiedS465
 from sharc.antenna.antenna_s580 import AntennaS580
@@ -1746,12 +1750,29 @@ def generate_mss_d2d(
                 1e6) + 30
         )
 
-       # Configure satellite positions in the StationManager
+        # Configure satellite positions in the StationManager
         x = mss_d2d_values["sat_x"]
         y = mss_d2d_values["sat_y"]
         z = mss_d2d_values["sat_z"]
         elev = mss_d2d_values["sat_antenna_elev"]
         azim = mss_d2d_values["sat_antenna_azim"]
+        global_ref = ENUReferenceFrame(
+            lat=coordinate_system.ref_lat,
+            lon=coordinate_system.ref_long,
+            alt=coordinate_system.ref_alt,
+        )
+        mss_d2d.geom.setup(
+            mss_d2d.num_stations, True,
+            global_ref,
+        )
+        mss_d2d.geom.set_local_reference_frame(
+            DWNReferenceFrame(
+                lat=mss_d2d_values["sat_lat"],
+                lon=mss_d2d_values["sat_lon"],
+                alt=mss_d2d_values["sat_alt"],
+            )
+        )
+
         mss_d2d.geom.set_global_coords(
             x, y, z,
             azim, elev,
@@ -1781,6 +1802,8 @@ def generate_mss_d2d(
             antenna_pattern = AntennaS1528(params.antenna.itu_r_s_1528)
         elif params.antenna.pattern == "ITU-R-S.1528-Taylor":
             antenna_pattern = AntennaS1528Taylor(params.antenna.itu_r_s_1528)
+        elif params.antenna.pattern == "ARRAY2":
+            pass
         elif params.antenna.pattern == "MSS Adjacent":
             antenna_pattern = AntennaMSSAdjacent(params.frequency)
         else:
@@ -1788,8 +1811,25 @@ def generate_mss_d2d(
                 f"generate_mss_ss: Invalid antenna type: {params.antenna.pattern}")
 
         for i in range(mss_d2d.num_stations):
+            if params.antenna.pattern == "ARRAY2":
+                antenna_pattern = AntennaArray(
+                    params.antenna.array,
+                    mss_d2d.geom.global2local.take(i)
+                )
+                antenna_pattern.add_beam(
+                    mss_d2d.geom.pointn_azim_global[i],
+                    90. - mss_d2d.geom.pointn_elev_global[i],
+                )
+                antenna_pattern.set_always_first_beam()
+
             mss_d2d.antenna[i] = antenna_pattern
 
+        if params.antenna.pattern == "ARRAY2":
+            mss_d2d.geom.set_local_coords(
+                azim=np.zeros_like(mss_d2d.geom.pointn_azim_global),
+                elev=np.zeros_like(mss_d2d.geom.pointn_elev_global),
+            )
+
         return mss_d2d  # Return the configured StationManager
 
     @staticmethod
@@ -1982,8 +2022,8 @@ def get_random_position(num_stas: int,
     parameters.imt.topology.sampling_from_spherical_grid.grid.grid_in_zone.circle.center_lon = ref_long
     parameters.imt.topology.sampling_from_spherical_grid.grid.grid_in_zone.circle.radius_km = 30 * 111
 
-    parameters.imt.topology.type = "SAMPLING_FROM_SPHERICAL_GRID"
-    # parameters.imt.topology.type = "MSS_DC"
+    # parameters.imt.topology.type = "SAMPLING_FROM_SPHERICAL_GRID"
+    parameters.imt.topology.type = "MSS_DC"
     parameters.imt.validate("station_factory_imt")
     # print(
     #     "parameters.imt.topology.sampling_from_spherical_grid.grid.lon_lat_grid.shape",
diff --git a/sharc/support/geometry.py b/sharc/support/geometry.py
index 94a93712..1fc9b548 100644
--- a/sharc/support/geometry.py
+++ b/sharc/support/geometry.py
@@ -354,8 +354,8 @@ def take(self, idx: int) -> "RigidTransform":
         # array shape structure for functional broadcasting
         # and batch computing
         return RigidTransform(
-            rot=self.rot[idx:idx+1],
-            t=self.t[idx:idx+1],
+            rot=self.rot[idx:idx + 1],
+            t=self.t[idx:idx + 1],
         )
 
     def apply_points(self, x: np.ndarray) -> np.ndarray:
@@ -453,6 +453,9 @@ def __init__(
         self._alt.flags.writeable = False
 
     def _compute_to_local(self):
+        return self._compute_to_enu()
+
+    def _compute_to_enu(self):
         lat = self._lat
         lon = self._lon
         alt = self._alt
@@ -475,6 +478,31 @@ def _compute_to_local(self):
         return ecef2local
 
 
+class DWNReferenceFrame(ENUReferenceFrame):
+    """
+    Defines DWN reference frame. x=Down, y=West, z=North
+    DWN is a custom reference frame for simplification of satellite's ref frame
+    useful for antenna usage.
+    """
+    ENU2DWN_ROT = scipy.spatial.transform.Rotation.from_matrix(
+        np.array([[
+            [0,  0, -1],
+            [-1,  0,  0],
+            [0,  1,  0],
+        ]])
+    )
+
+    def _compute_to_local(self):
+        ecef2enu = self._compute_to_enu()
+        # basis change
+        # ENU <> x: east, y: north, z: up
+        # DWN <> x: down, y: west, z: north
+        # x2 = -z1, y2 = -x1, z2 = y
+        enu2dwn = RigidTransform(self.ENU2DWN_ROT, np.zeros((1, 3)))
+
+        return ecef2enu.and_then(enu2dwn)
+
+
 @readonly_properties(
     "x_local", "y_local", "z_local",
     "pointn_azim_local", "pointn_elev_local",
@@ -501,7 +529,7 @@ def __init__(
         self,
         num_geometries,
         uses_local_coords=False,
-        global_cs: tuple[float, float, float] = None,
+        global_cs: ReferenceFrame = None,
     ):
         """
         Initialize a geometry object with a global coordinate system
@@ -554,6 +582,10 @@ def set_local_reference_frame(
         Sets local coord system references and prepares
         global<->local coordinate transformation
         """
+        if not self.uses_local_coords:
+            raise ValueError(
+                "Cannot set local reference frame when not using local coords"
+            )
         self._local_reference_frame = frame
 
         self._compute_global_local_transform()
@@ -926,14 +958,14 @@ def plot_geom(
     )
 
     bs_geom.set_local_reference_frame(
-        ENUReferenceFrame(
+        DWNReferenceFrame(
             lat=np.repeat(11, topology.num_base_stations),
             lon=np.repeat(-47, topology.num_base_stations),
             alt=np.repeat(1200, topology.num_base_stations),
         )
     )
     ue_geom.set_local_reference_frame(
-        ENUReferenceFrame(
+        DWNReferenceFrame(
             lat=np.repeat(11, num_ue * topology.num_base_stations),
             lon=np.repeat(-47, num_ue * topology.num_base_stations),
             alt=np.repeat(1200, num_ue * topology.num_base_stations),
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index a9f68c11..ee6fa261 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -1,8 +1,9 @@
 import unittest
 import numpy as np
 import numpy.testing as npt
+from sharc.support.sharc_utils import wrap2_180
 from sharc.support.geometry import (
-    SimulatorGeometry, ENUReferenceFrame, RigidTransform
+    SimulatorGeometry, DWNReferenceFrame, ENUReferenceFrame, RigidTransform
 )
 from sharc.satellite.ngso.constants import EARTH_RADIUS_M
 from copy import deepcopy
@@ -26,21 +27,25 @@ def random_rigid_transform(rng, N):
 
     return RigidTransform(Rotation.from_matrix(Q), t)
 
+
 def rot_identity(n=1):
     return Rotation.from_rotvec(np.zeros((n, 3)))
 
+
 def rot_z(angle_deg, n=1):
     return Rotation.from_rotvec(
         np.tile([0.0, 0.0, angle_deg], (n, 1)),
         degrees=True
     )
 
+
 def rot_x(angle_deg, n=1):
     return Rotation.from_rotvec(
         np.tile([angle_deg, 0.0, 0.0], (n, 1)),
         degrees=True
     )
 
+
 def rot_y(angle_deg, n=1):
     return Rotation.from_rotvec(
         np.tile([0.0, angle_deg, 0.0], (n, 1)),
@@ -69,11 +74,11 @@ def test_init_and_broadcasting(self):
 
             if Nt > 1:
                 with self.assertRaises(ValueError):
-                    RigidTransform(rot_identity(N+1), t)
+                    RigidTransform(rot_identity(N + 1), t)
 
             if Nrot > 1:
                 with self.assertRaises(ValueError):
-                    RigidTransform(rot, np.zeros((N+1, 3)))
+                    RigidTransform(rot, np.zeros((N + 1, 3)))
 
             # should not throw:
             tr = RigidTransform(rot, t)
@@ -84,7 +89,7 @@ def test_init_and_broadcasting(self):
             for fn in [tr.apply_points, tr.apply_vectors]:
                 if N > 1:
                     with self.assertRaises(AssertionError):
-                        fn(np.zeros((N+1, 3)))
+                        fn(np.zeros((N + 1, 3)))
 
                 for in_shp in [(1, 3), (N, 3)]:
                     res = fn(np.zeros(in_shp))
@@ -94,7 +99,7 @@ def test_init_and_broadcasting(self):
 
             for fn, in_shp in product(
                 [tr.apply_points_permutation, tr.apply_vectors_permutation],
-                [(1, 3), (N, 3), (N+1, 3)]
+                [(1, 3), (N, 3), (N + 1, 3)]
             ):
                 Nin = in_shp[0]
                 res = fn(np.zeros(in_shp))
@@ -338,6 +343,53 @@ def test_take_matches_permutation_vectors(self):
                 )
 
 
+class TestDWNReferenceFrame(unittest.TestCase):
+    def setUp(self):
+        self.lat = np.array([0.0])
+        self.lon = np.array([0.0])
+        self.alt = np.array([0.0])
+
+        self.enu = ENUReferenceFrame(
+            lat=self.lat, lon=self.lon, alt=self.alt
+        )
+        self.dwn = DWNReferenceFrame(
+            lat=self.lat, lon=self.lon, alt=self.alt
+        )
+
+    def test_enu_to_dwn_basis(self):
+        # ENU basis vectors
+        e = np.array([1.0, 0.0, 0.0])  # East
+        n = np.array([0.0, 1.0, 0.0])  # North
+        u = np.array([0.0, 0.0, 1.0])  # Up
+
+        # Transform ENU -> ECEF -> DWN
+        e_dwn = self.dwn.from_ecef.apply_vectors(
+            self.enu.to_ecef.apply_vectors(e)
+        )[0]
+        n_dwn = self.dwn.from_ecef.apply_vectors(
+            self.enu.to_ecef.apply_vectors(n)
+        )[0]
+        u_dwn = self.dwn.from_ecef.apply_vectors(
+            self.enu.to_ecef.apply_vectors(u)
+        )[0]
+
+        npt.assert_allclose(e_dwn, np.array([0.0, -1.0, 0.0]), atol=1e-4)
+        npt.assert_allclose(n_dwn, np.array([0.0, 0.0, 1.0]), atol=1e-4)
+        npt.assert_allclose(u_dwn, np.array([-1.0, 0.0, 0.0]), atol=1e-4)
+        npt.assert_allclose(
+            np.linalg.norm(e_dwn),
+            np.linalg.norm(e),
+        )
+        npt.assert_allclose(
+            np.linalg.norm(n_dwn),
+            np.linalg.norm(n),
+        )
+        npt.assert_allclose(
+            np.linalg.norm(u_dwn),
+            np.linalg.norm(u),
+        )
+
+
 class TestGeometry(unittest.TestCase):
     """Unit tests for the CoordinateSystem class and related coordinate transformations."""
 

From e957c90ca4c9ac2f1542e1ba5e9d1c54f4016d37 Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 18:31:02 -0300
Subject: [PATCH 06/10] update: better performance for antenna array
 implementation

---
 sharc/antenna/antenna_array.py | 78 +++++++++++++++++++++++++++-------
 1 file changed, 63 insertions(+), 15 deletions(-)

diff --git a/sharc/antenna/antenna_array.py b/sharc/antenna/antenna_array.py
index 18dea309..6c7f0dab 100644
--- a/sharc/antenna/antenna_array.py
+++ b/sharc/antenna/antenna_array.py
@@ -100,7 +100,7 @@ def _array_gain(
             beam_phi, beam_theta = self._to_local_coord(beam_phi, beam_theta)
 
         beam_etilt = beam_theta - 90.
-        beams_w_vec_row, beams_w_vec_col = self._weight_vector(
+        beams_w_vec_row, beams_w_vec_col = self._weight_vector_components(
             beam_phi, beam_etilt,
             self.par.n_rows, self.par.n_columns,
             self.par.element_vert_spacing,
@@ -108,7 +108,7 @@ def _array_gain(
         )
         w_vec_row, w_vec_col = beams_w_vec_row[beam_idxs], beams_w_vec_col[beam_idxs]
 
-        v_vec_row, v_vec_col = self._super_position_vector(
+        v_vec_row, v_vec_col = self._super_position_vector_components(
             phi, theta,
             self.par.n_rows, self.par.n_columns,
             self.par.element_vert_spacing,
@@ -127,12 +127,44 @@ def _array_gain(
 
         return g
 
+    @staticmethod
+    def _super_position_vector(
+        phi_tilt: np.ndarray,
+        theta_tilt: np.ndarray,
+        n_rows: int, n_cols: int,
+        dv: float, dh: float,
+    ) -> np.array:
+        vn, vm = AntennaArray._super_position_vector_components(
+            phi_tilt,
+            theta_tilt,
+            n_rows, n_cols,
+            dv, dh,
+        )
+
+        return vn[:, :, None] * vm[:, None, :]
+
     @staticmethod
     def _weight_vector(
         phi_tilt: np.ndarray,
         theta_tilt: np.ndarray,
         n_rows: int, n_cols: int,
         dv: float, dh: float,
+    ) -> np.array:
+        wn, wm = AntennaArray._weight_vector_components(
+            phi_tilt,
+            theta_tilt,
+            n_rows, n_cols,
+            dv, dh,
+        )
+
+        return wn[:, :, None] * wm[:, None, :]
+
+    @staticmethod
+    def _weight_vector_components(
+        phi_tilt: np.ndarray,
+        theta_tilt: np.ndarray,
+        n_rows: int, n_cols: int,
+        dv: float, dh: float,
     ) -> np.array:
         """
         Calculates super position vector.
@@ -162,19 +194,18 @@ def _weight_vector(
         m = np.arange(n_cols)[np.newaxis, :] + 1
 
         exp_arg_n = (n - 1) * dv * np.sin(r_theta)
-        exp_arg_m = (m - 1) * dh * np.cos(r_theta) * np.sin(r_phi)
+        exp_arg_m = - (m - 1) * dh * np.cos(r_theta) * np.sin(r_phi)
 
         w_vec_n = (1 / np.sqrt(n_rows * n_cols)) *\
             np.exp(2 * np.pi * 1.0j * exp_arg_n)
 
-        w_vec_m = (1 / np.sqrt(n_rows * n_cols)) *\
-            np.exp(2 * np.pi * 1.0j * exp_arg_m)
+        w_vec_m = np.exp(2 * np.pi * 1.0j * exp_arg_m)
 
         # shape (Na, Nr, Nc)
         return (w_vec_n, w_vec_m)
 
     @staticmethod
-    def _super_position_vector(
+    def _super_position_vector_components(
         phi: float, theta: float,
         n_rows: int, n_cols: int,
         dv: float, dh: float,
@@ -199,18 +230,35 @@ def _super_position_vector(
         A = dv * np.cos(np.deg2rad(theta))
         B = dh * np.sin(np.deg2rad(theta)) * np.sin(np.deg2rad(phi))
 
-        # indices
-        n = np.arange(n_rows)          # (Nr,)
-        m = np.arange(n_cols)          # (Nc,)
+        # instead of calculating exp for every row, there is a recursive
+        # relation that speeds this up. Small n_rows means that floating
+        # point error should not accumulate
+        # The relationship is: V(n) = V(n-1)V(2) | n > 2
+        # V(1) = 1., V(2) = exp(...)
+        row_phase = np.empty((len(theta), n_rows), dtype=np.complex128)
+        row_phase[:, 0] = 1.
 
-        # (Na, Nr)
-        row_phase = np.exp(
-            2j * np.pi * A[:, None] * n[None, :]
+        row_phase_term = np.exp(
+            2j * np.pi * A
         )
-        # (Na, Nc)
-        col_phase = np.exp(
-            2j * np.pi * B[:, None] * m[None, :]
+        row_phase[:, 1:] = row_phase_term[:, None]
+        # recursive relationship by cumulative product
+        row_phase = np.cumprod(row_phase, axis=-1)
+
+        # instead of calculating exp for every col, there is a recursive
+        # relation that speeds this up. Small n_cols means that floating
+        # point error should not accumulate
+        # The relationship is: V(n) = V(n-1)V(2) | n > 2
+        # V(1) = 1., V(2) = exp(...)
+        col_phase = np.empty((len(theta), n_cols), dtype=np.complex128)
+        col_phase[:, 0] = 1.
+
+        col_phase_term = np.exp(
+            2j * np.pi * B
         )
+        col_phase[:, 1:] = col_phase_term[:, None]
+        # recursive relationship by cumulative product
+        col_phase = np.cumprod(col_phase, axis=-1)
 
         return (row_phase, col_phase)
 

From 0dd5d9bf45901914955f30bd6c30390cf82dd3fd Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 22:57:24 -0300
Subject: [PATCH 07/10] update: antenna array comments, signatures and minor
 changes

---
 sharc/antenna/antenna_array.py | 150 ++++++++++++++++++++++++++++-----
 1 file changed, 128 insertions(+), 22 deletions(-)

diff --git a/sharc/antenna/antenna_array.py b/sharc/antenna/antenna_array.py
index 6c7f0dab..80c4352f 100644
--- a/sharc/antenna/antenna_array.py
+++ b/sharc/antenna/antenna_array.py
@@ -1,4 +1,6 @@
 """
+Optimized implementation of M.2101 antenna array.
+
 This antenna was created after the already existing antenna_beamforming_imt
 since that implementation is too slow for use with a lot of stations.
 
@@ -12,16 +14,34 @@
 from sharc.support.sharc_geom import polar_to_cartesian, cartesian_to_polar
 
 import numpy as np
+import typing
 
 
 class AntennaArray(Antenna):
-    # par: ParametersAntennaImt
+    """Implements M.2101 antenna array."""
 
     def __init__(
         self,
         par: ParametersAntennaImt,
         global2local_transform: RigidTransform = None,
     ):
+        """Constructs antenna array.
+
+        Parameters
+        ----------
+        par: ParametersAntennaImt
+            Antenna parameters. Partial support only.
+        global2local_transform: RigidTransform, optional
+            Transformation from global to local coordinate system. If None,
+            no transformation is applied.
+
+        Notes
+        -----
+        By partial support, it is meant that not all parameters
+        from ParametersAntennaImt are used in this implementation.
+        For example, normalization and subarray support are not
+        implemented.
+        """
         super().__init__()
         self.par = par
         self.always_first_beam = False
@@ -35,15 +55,33 @@ def __init__(
                 )
 
     def set_always_first_beam(self):
-        """
+        """Sets the antenna to always use the first beam.
+
         In case this is called, then calculate_gains will sum all beams
         contributions for each direction angle.
         """
         self.always_first_beam = True
 
     def calculate_gain(self, *args, **kwargs) -> np.array:
-        """
-        Calculates the antenan gain.
+        """Calculates the antenna gain.
+
+        Parameters
+        ----------
+        phi_vec : np.ndarray
+            Azimuth angles [degrees] in global coordinate system.
+        theta_vec : np.ndarray
+            Elevation angles [degrees] in global coordinate system.
+        co_channel : bool, optional
+            If True, co-channel interference is considered (default is True).
+        beams_l : np.ndarray, optional
+            Indices of beams to consider for each angle. If not provided,
+            all beams are considered. Also, if always_first_beam is set,
+            this parameter is ignored.
+
+        Returns
+        -------
+        np.ndarray
+            Antenna gain [dBi] for each direction.
         """
         phi_vec = np.atleast_1d(kwargs["phi_vec"])
         theta_vec = np.atleast_1d(kwargs["theta_vec"])
@@ -84,6 +122,16 @@ def calculate_gain(self, *args, **kwargs) -> np.array:
     def _element_gain(
         self, phi: np.ndarray, theta: np.ndarray
     ):
+        """
+        Calculates the element gain for given angles.
+
+        Parameters
+        ----------
+        phi : np.ndarray
+            Azimuth angles [degrees] in local coordinate system.
+        theta : np.ndarray
+            Elevation angles [degrees] in local coordinate system.
+        """
         return self._element_gain_dispatch(
             self.par, phi, theta,
         )
@@ -93,6 +141,15 @@ def _array_gain(
         phi: np.ndarray, theta: np.ndarray,
         beam_idxs: np.ndarray
     ):
+        """Calculates the array gain for given angles and beam indices.
+
+        Notes
+        -----
+        The mathematical formulation is based on M.2101, but formulation
+        has been optimized for computational efficiency. It considers
+        separability of the array factor into row and column components,
+        allowing for reduced memory bandwidth and faster runtime.
+        """
         if len(self.beams_list) == 0:
             beam_phi, beam_theta = phi, theta
         else:
@@ -106,7 +163,10 @@ def _array_gain(
             self.par.element_vert_spacing,
             self.par.element_horiz_spacing,
         )
-        w_vec_row, w_vec_col = beams_w_vec_row[beam_idxs], beams_w_vec_col[beam_idxs]
+        w_vec_row, w_vec_col = (
+            beams_w_vec_row[beam_idxs],
+            beams_w_vec_col[beam_idxs]
+        )
 
         v_vec_row, v_vec_col = self._super_position_vector_components(
             phi, theta,
@@ -125,7 +185,7 @@ def _array_gain(
             )**2
         )
 
-        return g
+        return np.maximum(g, self.par.minimum_array_gain)
 
     @staticmethod
     def _super_position_vector(
@@ -165,19 +225,21 @@ def _weight_vector_components(
         theta_tilt: np.ndarray,
         n_rows: int, n_cols: int,
         dv: float, dh: float,
-    ) -> np.array:
+    ) -> typing.Tuple[np.ndarray, np.ndarray]:
         """
         Calculates super position vector.
         Angles are in the local coordinate system.
 
         Parameters
         ----------
-            phi_tilt (float): electrical horizontal steering [degrees]
-            theta_tilt (float): electrical down-tilt steering [degrees]
+            phi_tilt: np.ndarray
+                electrical horizontal steering [degrees]
+            theta_tilt: np.ndarray
+                electrical down-tilt steering [degrees]
 
         Returns
         -------
-            w_vec (np.array, np.array):
+            w_vec: (np.ndarray, np.ndarray)
                 weighting vectors, first for rows, second for columns
         """
         # shape (Na, 1, 1)
@@ -201,7 +263,6 @@ def _weight_vector_components(
 
         w_vec_m = np.exp(2 * np.pi * 1.0j * exp_arg_m)
 
-        # shape (Na, Nr, Nc)
         return (w_vec_n, w_vec_m)
 
     @staticmethod
@@ -209,19 +270,27 @@ def _super_position_vector_components(
         phi: float, theta: float,
         n_rows: int, n_cols: int,
         dv: float, dh: float,
-    ) -> np.array:
+    ) -> typing.Tuple[np.ndarray, np.ndarray]:
         """
         Calculates super position vector.
         Angles are in the local coordinate system.
 
         Parameters
         ----------
-            theta (float): elevation angle [degrees]
-            phi (float): azimuth angle [degrees]
+            theta: float
+                elevation angle [degrees]
+            phi: float
+                azimuth angle [degrees]
 
         Returns
         -------
-            v_vec (np.array): superposition vector
+        v_vec: typing.Tuple[np.ndarray, np.ndarray]:
+            superposition vector components, first for rows, second for columns
+
+        Notes
+        -----
+        This implementation is optimized for computational efficiency,
+        using recursive relationships to avoid redundant calculations.
         """
         phi = np.atleast_1d(phi)
         theta = np.atleast_1d(theta)
@@ -243,7 +312,7 @@ def _super_position_vector_components(
         )
         row_phase[:, 1:] = row_phase_term[:, None]
         # recursive relationship by cumulative product
-        row_phase = np.cumprod(row_phase, axis=-1)
+        np.cumprod(row_phase, axis=-1, out=row_phase)
 
         # instead of calculating exp for every col, there is a recursive
         # relation that speeds this up. Small n_cols means that floating
@@ -258,12 +327,24 @@ def _super_position_vector_components(
         )
         col_phase[:, 1:] = col_phase_term[:, None]
         # recursive relationship by cumulative product
-        col_phase = np.cumprod(col_phase, axis=-1)
+        np.cumprod(col_phase, axis=-1, out=col_phase)
 
         return (row_phase, col_phase)
 
     @staticmethod
     def _element_gain_dispatch(par: ParametersAntennaImt, phi, theta):
+        """
+        Dispatch to the correct element gain calculation method.
+
+        Parameters
+        ----------
+        par: ParametersAntennaImt
+            Antenna parameters.
+        phi: np.ndarray
+            Azimuth angles [degrees] in local coordinate system.
+        theta: np.ndarray
+            Elevation angles [degrees] in local coordinate system.
+        """
         if par.element_pattern == "M2101":
             return AntennaArray._calculate_m2101_element_gain(
                 phi, theta,
@@ -284,8 +365,7 @@ def _calculate_m2101_element_gain(
         g_max: np.ndarray, sla_v: np.ndarray, am: np.ndarray,
         multiplication_factor: np.ndarray = 12
     ):
-        """Calculates and returns element gain as described in M.2101
-        """
+        """Calculates and returns element gain as described in M.2101."""
         g_horizontal = -1.0 * np.minimum(
             multiplication_factor * (phi / phi_3db)**2, am
         )
@@ -301,6 +381,31 @@ def _calculate_m2101_element_gain(
         return g_max - np.minimum(att, am)
 
     def _to_local_coord(self, phi, theta):
+        """
+        Transform angles from global to local coordinate system.
+
+        Parameters
+        ----------
+        phi: np.ndarray
+            Azimuth angles [degrees] in global coordinate system.
+        theta: np.ndarray
+            Elevation angles [degrees] in global coordinate system.
+
+        Returns
+        -------
+        phi: np.ndarray
+            Azimuth angles [degrees] in local coordinate system.
+        theta: np.ndarray
+            Elevation angles [degrees] in local coordinate system.
+
+        Notes
+        -----
+        The transformation is done by converting to Cartesian coordinates,
+        applying the transformation defined on construction,
+        and converting back to spherical coordinates. It is assumed that
+        theta is defined with z axis as reference, and phi with x axis as reference
+        and increasing towards y axis.
+        """
         if self.global2local_transform is None:
             return np.array(phi), np.array(theta)
 
@@ -324,12 +429,13 @@ def add_beam(self, phi_etilt: float, theta_etilt: float):
 
         Parameters
         ----------
-            phi_etilt (float): azimuth electrical tilt angle [degrees]
-            theta_etilt (float): elevation electrical tilt angle [degrees]
+            phi_etilt: float
+                azimuth electrical tilt angle [degrees]
+            theta_etilt: float
+                elevation electrical tilt angle [degrees]
         """
         phi_etilt, theta_etilt = np.atleast_1d(phi_etilt), np.atleast_1d(theta_etilt)
 
-    # def add_beam_in_local_coords(self):
         self.beams_list.append(
             (np.ndarray.item(phi_etilt), np.ndarray.item(theta_etilt)),
         )

From f714def774a89311f792f2330242340210a62501 Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Mon, 2 Feb 2026 22:58:11 -0300
Subject: [PATCH 08/10] fix: missing comments

---
 sharc/support/geometry.py | 11 ++++++++---
 tests/test_geometry.py    | 18 ++++++++++++++++++
 2 files changed, 26 insertions(+), 3 deletions(-)

diff --git a/sharc/support/geometry.py b/sharc/support/geometry.py
index 1fc9b548..c6918019 100644
--- a/sharc/support/geometry.py
+++ b/sharc/support/geometry.py
@@ -333,11 +333,13 @@ def __post_init__(self):
 
     @property
     def N(self):
+        """Number of transforms represented"""
         n_rot = len(self.rot)
         n_t = self.t.shape[0]
         return max(n_rot, n_t)
 
     def inv(self) -> "RigidTransform":
+        """Returns the inverse transform"""
         rot_inv = self.rot.inv()
         t_inv = -rot_inv.apply(self.t)
         return RigidTransform(rot_inv, t_inv)
@@ -350,6 +352,7 @@ def and_then(self, other: "RigidTransform") -> "RigidTransform":
         )
 
     def take(self, idx: int) -> "RigidTransform":
+        """Takes the idx-th transform from the batch"""
         # NOTE: we slice instead of taking indice to maintain
         # array shape structure for functional broadcasting
         # and batch computing
@@ -405,10 +408,12 @@ class ReferenceFrame(ABC):
 
     @property
     def from_ecef(self) -> RigidTransform:
+        """Returns the transform from ECEF to the local coordinate system."""
         return self._from_ecef
 
     @property
     def to_ecef(self) -> RigidTransform:
+        """Returns the transform from the local coordinate system to ECEF."""
         return self._to_ecef
 
 
@@ -486,9 +491,9 @@ class DWNReferenceFrame(ENUReferenceFrame):
     """
     ENU2DWN_ROT = scipy.spatial.transform.Rotation.from_matrix(
         np.array([[
-            [0,  0, -1],
-            [-1,  0,  0],
-            [0,  1,  0],
+            [0, 0, -1],
+            [-1, 0, 0],
+            [0, 1, 0],
         ]])
     )
 
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index ee6fa261..5eb3c881 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -29,10 +29,12 @@ def random_rigid_transform(rng, N):
 
 
 def rot_identity(n=1):
+    """Returns identity rotation of batch size n."""
     return Rotation.from_rotvec(np.zeros((n, 3)))
 
 
 def rot_z(angle_deg, n=1):
+    """Returns rotation about Z axis by angle_deg degrees, batch size n."""
     return Rotation.from_rotvec(
         np.tile([0.0, 0.0, angle_deg], (n, 1)),
         degrees=True
@@ -40,6 +42,7 @@ def rot_z(angle_deg, n=1):
 
 
 def rot_x(angle_deg, n=1):
+    """Returns rotation about X axis by angle_deg degrees, batch size n."""
     return Rotation.from_rotvec(
         np.tile([angle_deg, 0.0, 0.0], (n, 1)),
         degrees=True
@@ -47,6 +50,7 @@ def rot_x(angle_deg, n=1):
 
 
 def rot_y(angle_deg, n=1):
+    """Returns rotation about Y axis by angle_deg degrees, batch size n."""
     return Rotation.from_rotvec(
         np.tile([0.0, angle_deg, 0.0], (n, 1)),
         degrees=True
@@ -54,10 +58,14 @@ def rot_y(angle_deg, n=1):
 
 
 class TestRigidTransform(unittest.TestCase):
+    """Unit tests for RigidTransform class."""
+
     def setUp(self):
+        """Set up test fixtures for RigidTransform tests."""
         pass
 
     def test_init_and_broadcasting(self):
+        """Test RigidTransform initialization and broadcasting behavior."""
         for rot_shp, t_shp in product(
             [(1, 3), (4, 3)],
             [(1, 3), (4, 3)],
@@ -108,6 +116,7 @@ def test_init_and_broadcasting(self):
                 )
 
     def test_simple_transformations(self):
+        """Test RigidTransform simple transformations and their combinations."""
         eps = 1e-4
 
         ux = np.array([1., 0., 0.])
@@ -250,6 +259,7 @@ def test_simple_transformations(self):
         )
 
     def test_permutation_points_equivalence(self):
+        """Test equivalence of permutation and non-permutation point applications."""
         rng = np.random.default_rng(0)
 
         for n in range(1, 10):
@@ -268,6 +278,7 @@ def test_permutation_points_equivalence(self):
             )
 
     def test_permutation_vectors_equivalence(self):
+        """Test equivalence of permutation and non-permutation vector applications."""
         rng = np.random.default_rng(0)
 
         for n in range(1, 10):
@@ -286,6 +297,7 @@ def test_permutation_vectors_equivalence(self):
             )
 
     def test_take_commutes_with_apply_points(self):
+        """Test that 'take' method commutes with apply_points."""
         rng = np.random.default_rng(0)
 
         for n in range(1, 10):
@@ -303,6 +315,7 @@ def test_take_commutes_with_apply_points(self):
                 )
 
     def test_take_matches_permutation_points(self):
+        """Test that 'take' method matches permutation point applications."""
         rng = np.random.default_rng(1)
 
         for n in [2, 5]:
@@ -323,6 +336,7 @@ def test_take_matches_permutation_points(self):
                 )
 
     def test_take_matches_permutation_vectors(self):
+        """Test that 'take' method matches permutation vector applications."""
         rng = np.random.default_rng(1)
 
         for n in [2, 5]:
@@ -344,7 +358,10 @@ def test_take_matches_permutation_vectors(self):
 
 
 class TestDWNReferenceFrame(unittest.TestCase):
+    """Unit tests for DWNReferenceFrame class."""
+
     def setUp(self):
+        """Set up test fixtures for DWNReferenceFrame tests."""
         self.lat = np.array([0.0])
         self.lon = np.array([0.0])
         self.alt = np.array([0.0])
@@ -357,6 +374,7 @@ def setUp(self):
         )
 
     def test_enu_to_dwn_basis(self):
+        """Test transformation of basis vectors from ENU to DWN."""
         # ENU basis vectors
         e = np.array([1.0, 0.0, 0.0])  # East
         n = np.array([0.0, 1.0, 0.0])  # North

From e7fd325a92443111cdbea267d3d534bd75711eea Mon Sep 17 00:00:00 2001
From: "Artur P. P." <89547366+artistrea@users.noreply.github.com>
Date: Tue, 3 Feb 2026 00:25:07 -0300
Subject: [PATCH 09/10] refactor: comments and docstring for better
 documentation

Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com>
---
 sharc/antenna/antenna_array.py | 6 +++---
 sharc/station_factory.py       | 2 +-
 tests/test_antenna_array.py    | 7 ++-----
 tests/test_geometry.py         | 2 +-
 4 files changed, 7 insertions(+), 10 deletions(-)

diff --git a/sharc/antenna/antenna_array.py b/sharc/antenna/antenna_array.py
index 80c4352f..b639cfc3 100644
--- a/sharc/antenna/antenna_array.py
+++ b/sharc/antenna/antenna_array.py
@@ -57,8 +57,8 @@ def __init__(
     def set_always_first_beam(self):
         """Sets the antenna to always use the first beam.
 
-        In case this is called, then calculate_gains will sum all beams
-        contributions for each direction angle.
+        When this flag is set, :meth:`calculate_gain` ignores any ``beams_l``
+        argument and selects the first beam (index 0) for all direction angles.
         """
         self.always_first_beam = True
 
@@ -227,7 +227,7 @@ def _weight_vector_components(
         dv: float, dh: float,
     ) -> typing.Tuple[np.ndarray, np.ndarray]:
         """
-        Calculates super position vector.
+        Calculates the complex weight vectors for beamforming.
         Angles are in the local coordinate system.
 
         Parameters
diff --git a/sharc/station_factory.py b/sharc/station_factory.py
index a6666863..322e691f 100644
--- a/sharc/station_factory.py
+++ b/sharc/station_factory.py
@@ -11,7 +11,7 @@
 import math
 
 from sharc.support.geometry import (
-    DWNReferenceFrame, SimulatorGeometry, ENUReferenceFrame
+    DWNReferenceFrame, ENUReferenceFrame
 )
 from sharc.support.enumerations import StationType
 from sharc.parameters.parameters import Parameters
diff --git a/tests/test_antenna_array.py b/tests/test_antenna_array.py
index adaf1291..8a7c4cc0 100644
--- a/tests/test_antenna_array.py
+++ b/tests/test_antenna_array.py
@@ -64,7 +64,7 @@ def setUp(self):
 
     def test_element_gain(self):
         """Testing element gain calculations"""
-        # self.antenna1._calculate_m2101_element_gain()
+
         """Test M.2101 horizontal pattern calculation for various phi values."""
         # phi = 0 results in zero gain
         phi = np.array([0., 120., 150.])
@@ -355,7 +355,4 @@ def test_calculate_gain(self):
 
 if __name__ == '__main__':
     unittest.main()
-#
-#    suite = unittest.TestSuite()
-#    suite.addTest(AntennaArrayTest('test_calculate_gain'))
-#    unittest.TextTestRunner().run(suite)
+
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index 5eb3c881..a25fa33c 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -1,7 +1,7 @@
 import unittest
 import numpy as np
 import numpy.testing as npt
-from sharc.support.sharc_utils import wrap2_180
+
 from sharc.support.geometry import (
     SimulatorGeometry, DWNReferenceFrame, ENUReferenceFrame, RigidTransform
 )

From 04685253c649ec14ae90af0263f01f00a7ebea70 Mon Sep 17 00:00:00 2001
From: artistrea <apiravesi@gmail.com>
Date: Tue, 3 Feb 2026 08:46:31 -0300
Subject: [PATCH 10/10] fix: linting errors

---
 sharc/support/geometry.py   | 76 ++++++++++++++++++++++++++++++++++---
 tests/test_antenna_array.py |  1 -
 tests/test_geometry.py      |  6 ---
 3 files changed, 71 insertions(+), 12 deletions(-)

diff --git a/sharc/support/geometry.py b/sharc/support/geometry.py
index c6918019..98895f57 100644
--- a/sharc/support/geometry.py
+++ b/sharc/support/geometry.py
@@ -318,6 +318,7 @@ def __post_init__(self):
         if not isinstance(self.rot, scipy.spatial.transform.Rotation):
             raise ValueError("rot must be a scipy Rotation")
 
+        # NOTE: this throws if the rotation isn't defined as a batch rotation
         n_rot = len(self.rot)
         n_t = self.t.shape[0]
 
@@ -362,17 +363,63 @@ def take(self, idx: int) -> "RigidTransform":
         )
 
     def apply_points(self, x: np.ndarray) -> np.ndarray:
-        """Applies rotations and translations to the points"""
+        """Applies transform to the points.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Array of input points of shape (N, 3). Each row is a 3D point in
+            Cartesian coordinates.
+
+        Returns
+        -------
+        np.ndarray
+            Array of transformed points of shape (N, 3), where ``N`` is the
+            number of transforms represented by this ``RigidTransform``.
+        """
         return self.apply_vectors(x) + self.t
 
     def apply_points_permutation(self, x: np.ndarray) -> np.ndarray:
-        """Applies rotation and translation to the point"""
+        """Applies transform to points with permutation broadcasting.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Array of input points of shape (M, 3). Each row is a 3D point in
+            Cartesian coordinates.
+
+        Returns
+        -------
+        np.ndarray
+            Array of transformed points of shape (N, M, 3), where ``N`` is the
+            number of transforms represented by this ``RigidTransform``. The
+            slice ``result[n, :, :]`` contains all ``M`` input points
+            transformed by the ``n``-th rigid transform.
+        """
         t = self.t                     # (N,3)
 
         return self.apply_vectors_permutation(x) + t[:, None, :]
 
     def apply_vectors(self, v: np.ndarray) -> np.ndarray:
-        """Applies only rotations, considering the vectors as a pointing vec"""
+        """Applies transform rotation-only to the points.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Array of input points of shape (N, 3). Each row is a 3D point in
+            Cartesian coordinates.
+
+        Returns
+        -------
+        np.ndarray
+            Array of rotated points of shape (N, 3), where ``N`` is the
+            number of transforms represented by this ``RigidTransform``.
+
+        Note
+        ----
+        This only applies rotation, it ignores translation. It is useful for
+        transforming pointing vectors.
+        """
         v = np.atleast_2d(v)
         assert v.ndim == 2 and v.shape[1] == 3
         M = v.shape[0]
@@ -381,8 +428,27 @@ def apply_vectors(self, v: np.ndarray) -> np.ndarray:
         return self.rot.apply(v) + np.zeros_like(self.t)
 
     def apply_vectors_permutation(self, v: np.ndarray) -> np.ndarray:
-        """Applies rotation and translation to the point"""
-        # TODO: change impl
+        """Applies rotation-only to points with permutation broadcasting.
+
+        Parameters
+        ----------
+        x : np.ndarray
+            Array of input points of shape (M, 3). Each row is a 3D point in
+            Cartesian coordinates.
+
+        Returns
+        -------
+        np.ndarray
+            Array of transformed points of shape (N, M, 3), where ``N`` is the
+            number of transforms represented by this ``RigidTransform``. The
+            slice ``result[n, :, :]`` contains all ``M`` input points
+            transformed by the ``n``-th rigid transform rotation.
+
+        Note
+        ----
+        This only applies rotation, it ignores translation. It is useful for
+        transforming pointing vectors.
+        """
         v = np.atleast_2d(v)
         assert v.ndim == 2 and v.shape[1] == 3
 
diff --git a/tests/test_antenna_array.py b/tests/test_antenna_array.py
index 8a7c4cc0..686e0557 100644
--- a/tests/test_antenna_array.py
+++ b/tests/test_antenna_array.py
@@ -355,4 +355,3 @@ def test_calculate_gain(self):
 
 if __name__ == '__main__':
     unittest.main()
-
diff --git a/tests/test_geometry.py b/tests/test_geometry.py
index a25fa33c..9144eca5 100644
--- a/tests/test_geometry.py
+++ b/tests/test_geometry.py
@@ -425,9 +425,6 @@ def _test_expected_geom(
             np.stack((expect_local["x"], expect_local["y"], expect_local["z"])),
             atol=0.001
         )
-        # npt.assert_allclose(geom.x_local, expect_local["x"], atol=0.001)
-        # npt.assert_allclose(geom.y_local, expect_local["y"], atol=0.001)
-        # npt.assert_allclose(geom.z_local, expect_local["z"], atol=0.001)
         npt.assert_allclose(geom.pointn_elev_local, expect_local["elev"], atol=0.001)
         if only_local_azim is None:
             npt.assert_allclose(geom.pointn_azim_local, expect_local["azim"], atol=0.001)
@@ -444,9 +441,6 @@ def _test_expected_geom(
             np.stack((expect_global["x"], expect_global["y"], expect_global["z"])),
             atol=0.001
         )
-        # npt.assert_allclose(geom.x_global, expect_global["x"], atol=0.001)
-        # npt.assert_allclose(geom.y_global, expect_global["y"], atol=0.001)
-        # npt.assert_allclose(geom.z_global, expect_global["z"], atol=0.001)
         npt.assert_allclose(geom.pointn_elev_global, expect_global["elev"], atol=0.001)
         if only_global_azim is None:
             npt.assert_allclose(