From 3d1a4ce8d3e72da424adf47467cd948b6713d6bc Mon Sep 17 00:00:00 2001
From: Jakob Elias Wagner <jakob.e.wagner@tum.de>
Date: Tue, 23 Jan 2024 22:16:16 +0100
Subject: [PATCH] change bounding box input shape

---
 src/nos/utility/bounding_box.py               | 32 ++++---------------
 .../helmholtz/solver/test_adiabatic_layer.py  |  6 ++--
 test/utility/test_bounding_box.py             | 18 +++++------
 3 files changed, 19 insertions(+), 37 deletions(-)

diff --git a/src/nos/utility/bounding_box.py b/src/nos/utility/bounding_box.py
index d3d3f1e7..8818e322 100644
--- a/src/nos/utility/bounding_box.py
+++ b/src/nos/utility/bounding_box.py
@@ -1,5 +1,4 @@
 import dataclasses
-from typing import Tuple
 
 import numpy as np
 
@@ -20,35 +19,18 @@ def __post_init__(self):
         self.min = np.array([self.x_min, self.y_min, 0.0])  # 3d to be able to handle 3d values
         self.max = np.array([self.x_max, self.y_max, 0.0])
 
-    def check_input_array(self, x: np.array) -> Tuple[int, np.array]:
-        """Checks that the dimensions align with the underling implementations.
-
-        Args:
-            x: input array
-
-        Returns:
-            dimensions of the input array and the correctly shaped array.
-        """
-        try:
-            n_dim = x.shape[1]
-        except IndexError:
-            n_dim = x.shape[0]
-            x = x.reshape((-1, n_dim))
-        return n_dim, x
-
     def inside(self, x: np.array) -> np.ndarray:
         """Method to tell which points from the input array are inside the bounding box.
 
         Args:
-            x: Array that should be probed.
+            x: (n_dim, n) array that should be probed.
 
         Returns:
             indices to all values in the array which are inside the bounding box.
         """
-        n_dim, x = self.check_input_array(x)
 
-        x_inside = (x[:, 0] >= self.x_min) & (x[:, 0] <= self.x_max)
-        y_inside = (x[:, 1] >= self.y_min) & (x[:, 1] <= self.y_max)
+        x_inside = (x[0, :] >= self.x_min) & (x[0, :] <= self.x_max)
+        y_inside = (x[1, :] >= self.y_min) & (x[1, :] <= self.y_max)
 
         return np.where(x_inside & y_inside)[0]
 
@@ -59,16 +41,16 @@ def distance(self, x: np.array) -> np.array:
         their distance to the box is zero. The Frobenius norm is used to calculate the distance.
 
         Args:
-            x: (n, n_dim) array of coordinates.
+            x: (n_dim, n) array of coordinates.
 
         Returns:
             array of distances.
         """
-        n_dim, x = self.check_input_array(x)
+        n_dim = x.shape[0]
 
-        clamped = np.maximum(self.min[:n_dim], np.minimum(x, self.max[:n_dim]))
+        clamped = np.maximum(self.min[:n_dim, np.newaxis], np.minimum(x, self.max[:n_dim, np.newaxis]))
 
         # Compute the distance from the clamped points to the original points
-        dist = np.linalg.norm(clamped - x, axis=1)
+        dist = np.linalg.norm(clamped - x, axis=0)
 
         return dist
diff --git a/test/data/helmholtz/solver/test_adiabatic_layer.py b/test/data/helmholtz/solver/test_adiabatic_layer.py
index 2b2dd6f4..978acb47 100644
--- a/test/data/helmholtz/solver/test_adiabatic_layer.py
+++ b/test/data/helmholtz/solver/test_adiabatic_layer.py
@@ -33,7 +33,7 @@ def all_sides_absorber():
 
 
 def test_inside_zero(all_sides_absorber):
-    x = np.random.random((1000, 2)) @ np.array([[3.0, 0.0], [0.0, 1.0]]) - np.array([0.5, 0])  # all inside
+    x = np.array([[3.0, 0.0], [0.0, 1.0]]) @ np.random.random((2, 1000)) - np.array([[0.5, 0]]).T  # all inside
     assert np.allclose(all_sides_absorber.eval(x), 0)
 
 
@@ -45,12 +45,12 @@ def test_outside_not_zero(all_sides_absorber):
             np.concatenate([np.random.random((500,)) - 1.1, np.random.random((500,)) + 2.1]),
         ],
         axis=1,
-    )
+    ).T
     assert all(all_sides_absorber.eval(x) != 0)
 
 
 def test_correct_value(all_sides_absorber):
-    x = np.array([[4.0, 1.0], [-1, -1]])
+    x = np.array([[4.0, 1.0], [-1, -1]]).T
 
     sigma_0 = -3 * np.log(10) / 2.0
     sol = sigma_0 * 1j * np.array([1.5, np.sqrt(1.25)]) ** 2
diff --git a/test/utility/test_bounding_box.py b/test/utility/test_bounding_box.py
index 213e20bf..b2e2f7f1 100644
--- a/test/utility/test_bounding_box.py
+++ b/test/utility/test_bounding_box.py
@@ -5,21 +5,21 @@
 
 def test_inside_with_points_inside():
     bbox = BoundingBox2D(x_min=0, y_min=0, x_max=10, y_max=10)
-    points = np.array([[1, 1], [5, 5], [9, 9]])
+    points = np.array([[1, 1], [5, 5], [9, 9]]).T
     inside_indices = bbox.inside(points)
     assert np.array_equal(inside_indices, np.array([0, 1, 2]))
 
 
 def test_inside_with_points_outside():
     bbox = BoundingBox2D(x_min=0, y_min=0, x_max=10, y_max=10)
-    points = np.array([[-1, -1], [11, 11], [20, 20]])
+    points = np.array([[-1, -1], [11, 11], [20, 20]]).T
     inside_indices = bbox.inside(points)
     assert len(inside_indices) == 0
 
 
 def test_inside_with_points_on_edges():
     bbox = BoundingBox2D(x_min=0, y_min=0, x_max=10, y_max=10)
-    points = np.array([[0, 0], [10, 10], [0, 10], [10, 0]])
+    points = np.array([[0, 0], [10, 10], [0, 10], [10, 0]]).T
     inside_indices = bbox.inside(points)
     assert np.array_equal(inside_indices, np.array([0, 1, 2, 3]))
 
@@ -33,7 +33,7 @@ def test_distance_greater_zero():
         dims.sort()
         bbox = BoundingBox2D(*list(dims))
 
-        x = 2 * np.random.random((sample_size, 2)) - 1
+        x = 2 * np.random.random((2, sample_size)) - 1
 
         dist = bbox.distance(x)
         assert all(dist >= 0)
@@ -44,9 +44,9 @@ def test_distance_inside():
     test_size = 10
     sample_size = 1000
     for _ in range(test_size):
-        x = 2 * np.random.random((sample_size, 2))
-        max_vals = np.max(x, axis=0)
-        min_vals = np.min(x, axis=0)
+        x = 2 * np.random.random((2, sample_size))
+        max_vals = np.max(x, axis=1)
+        min_vals = np.min(x, axis=1)
         bbox = BoundingBox2D(*list(min_vals), *list(max_vals))
 
         dist = bbox.distance(x)
@@ -62,7 +62,7 @@ def test_distance_outside():
         dims.sort()
         bbox = BoundingBox2D(*list(dims))
 
-        x = np.random.random((sample_size, 2)) - 2.0
+        x = np.random.random((2, sample_size)) - 2.0
 
         dist = bbox.distance(x)
         assert all(dist > 0)
@@ -86,7 +86,7 @@ def test_distance_array():
             [0, -4],
             [3, 5],
         ]
-    )
+    ).T
     correct = np.array([0, 0, 0, 0, 1, np.sqrt(2), 2, np.sqrt(8), 1, np.sqrt(2), 2, np.sqrt(10)])
     print(bbox.distance(x))
     assert np.allclose(bbox.distance(x), correct)